wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' of https://github.com/williamFalcon/pytorch-lightning

Browse Source
This commit is contained in:
William Falcon
2020-01-21 16:53:40 -05:00
parent 3513cb4df9 432a0bcd06
commit c5d4b87375
2 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pl_examples/full_examples/imagenet/imagenet_example.py
+9 -6
View File
@@ -34,9 +34,12 @@ class ImageNetLightningModel(pl.LightningModule):
self.hparams = hparams
self.model = models.__dict__[self.hparams.arch](pretrained=self.hparams.pretrained)
def forward(self, x):
return self.model(x)
def training_step(self, batch, batch_idx):
images, target = batch
output = self.model(images)
output = self.forward(images)
loss_val = F.cross_entropy(output, target)
acc1, acc5 = self.__accuracy(output, target, topk=(1, 5))
@@ -59,7 +62,7 @@ class ImageNetLightningModel(pl.LightningModule):
def validation_step(self, batch, batch_idx):
images, target = batch
output = self.model(images)
output = self.forward(images)
loss_val = F.cross_entropy(output, target)
acc1, acc5 = self.__accuracy(output, target, topk=(1, 5))
@@ -132,7 +135,7 @@ class ImageNetLightningModel(pl.LightningModule):
std=[0.229, 0.224, 0.225],
)
train_dir = os.path.join(self.hparams.data, 'train')
train_dir = os.path.join(self.hparams.data_path, 'train')
train_dataset = datasets.ImageFolder(
train_dir,
transforms.Compose([
@@ -162,7 +165,7 @@ class ImageNetLightningModel(pl.LightningModule):
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
)
val_dir = os.path.join(self.hparams.data, 'val')
val_dir = os.path.join(self.hparams.data_path, 'val')
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(val_dir, transforms.Compose([
transforms.Resize(256),
@@ -185,7 +188,7 @@ class ImageNetLightningModel(pl.LightningModule):
' (default: resnet18)')
parser.add_argument('--epochs', default=90, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--seed', type=int, default=None,
parser.add_argument('--seed', type=int, default=42,
help='seed for initializing training. ')
parser.add_argument('-b', '--batch-size', default=256, type=int,
metavar='N',
@@ -214,7 +217,7 @@ def get_args():
help='how many gpus')
parent_parser.add_argument('--distributed-backend', type=str, default='dp', choices=('dp', 'ddp', 'ddp2'),
help='supports three options dp, ddp, ddp2')
parent_parser.add_argument('--use-16bit', dest='use-16bit', action='store_true',
parent_parser.add_argument('--use-16bit', dest='use_16bit', action='store_true',
help='if true uses 16 bit precision')
parent_parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
help='evaluate model on validation set')
pytorch_lightning/trainer/training_loop.py
+1
View File
@@ -353,6 +353,7 @@ class TrainerTrainLoopMixin(ABC):
stop = should_stop and met_min_epochs
if stop:
self.main_progress_bar.close()
model.on_train_end()
return
self.main_progress_bar.close()