Repository: pytorch/vision

Issue:
Resume in the Segmentation example
## 🐛 Bug

The segmentation training [script](https://github.com/pytorch/vision/blob/7b60f4db9707d7afdbb87fd4e8ef6906ca014720/references/segmentation/train.py) doesn't seem to allow correct resuming. I think [lines 131-133](https://github.com/pytorch/vision/blob/7b60f4db9707d7afdbb87fd4e8ef6906ca014720/references/segmentation/train.py#L131) should be changed to:

```
    start_epoch = 0
    if args.resume:
        checkpoint = torch.load(args.resume, map_location='cpu')
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        start_epoch = checkpoint['epoch'] + 1
```

Then, [line 161](https://github.com/pytorch/vision/blob/7b60f4db9707d7afdbb87fd4e8ef6906ca014720/references/segmentation/train.py#L161) and [lines 167-173](https://github.com/pytorch/vision/blob/7b60f4db9707d7afdbb87fd4e8ef6906ca014720/references/segmentation/train.py#L167) should be changed respectively to:

```
    for epoch in range(start_epoch, args.epochs):
```

and

```
utils.save_on_master(
            {
                'model': model_without_ddp.state_dict(),
                'optimizer': optimizer.state_dict(),
                'lr_scheduler': lr_scheduler.state_dict(),
                'epoch': epoch
            },
```

Besides, is there a good reason to set [`batch_size=1`](https://github.com/pytorch/vision/blob/7b60f4db9707d7afdbb87fd4e8ef6906ca014720/references/segmentation/train.py#L120) (instead of `args.batch_size` or even `2*args.batch_size`) for `data_loader_test`?

Thanks.



Discussion:
> Besides, is there a good reason to set `batch_size=1` for `data_loader_test`?

As I understand this is a hack to prevent the test data loader from having duplicates as `torch.utils.data.distributed.DistributedSampler` always fills the dataset so that the total number of samples is divisible by `num_gpus*batch_size`. However, if the number of validation samples is not a multiple of `num_gpus` then the results will not be correct even if `batch_size=1`. 
The point you mention about resuming looks good to me. Do you want to open a pull request to address these? Are there other training script you noticed that had that?

Removing 'args' from the checkpoint is BC breaking though. We could decide not to load it when loading the checkpoint.

Resolving patch:
diff --git a/references/segmentation/train.py b/references/segmentation/train.py
--- a/references/segmentation/train.py
+++ b/references/segmentation/train.py
@@ -128,10 +128,6 @@ def main(args):
     if args.distributed:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
 
-    if args.resume:
-        checkpoint = torch.load(args.resume, map_location='cpu')
-        model.load_state_dict(checkpoint['model'])
-
     model_without_ddp = model
     if args.distributed:
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
@@ -157,8 +153,15 @@ def main(args):
         optimizer,
         lambda x: (1 - x / (len(data_loader) * args.epochs)) ** 0.9)
 
+    if args.resume:
+        checkpoint = torch.load(args.resume, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        args.start_epoch = checkpoint['epoch'] + 1
+
     start_time = time.time()
-    for epoch in range(args.epochs):
+    for epoch in range(args.start_epoch, args.epochs):
         if args.distributed:
             train_sampler.set_epoch(epoch)
         train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, args.print_freq)
@@ -168,6 +171,7 @@ def main(args):
             {
                 'model': model_without_ddp.state_dict(),
                 'optimizer': optimizer.state_dict(),
+                'lr_scheduler': lr_scheduler.state_dict(),
                 'epoch': epoch,
                 'args': args
             },
@@ -201,6 +205,8 @@ def parse_args():
     parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
     parser.add_argument('--output-dir', default='.', help='path where to save')
     parser.add_argument('--resume', default='', help='resume from checkpoint')
+    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
     parser.add_argument(
         "--test-only",
         dest="test_only",


Test patch: