We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `references/segmentation/train.py`
Content:
```
1 import datetime
2 import os
3 import time
4 
5 import torch
6 import torch.utils.data
7 from torch import nn
8 import torchvision
9 
10 from coco_utils import get_coco
11 import transforms as T
12 import utils
13 
14 
15 def get_dataset(name, image_set, transform):
16     def sbd(*args, **kwargs):
17         return torchvision.datasets.SBDataset(*args, mode='segmentation', **kwargs)
18     paths = {
19         "voc": ('/datasets01/VOC/060817/', torchvision.datasets.VOCSegmentation, 21),
20         "voc_aug": ('/datasets01/SBDD/072318/', sbd, 21),
21         "coco": ('/datasets01/COCO/022719/', get_coco, 21)
22     }
23     p, ds_fn, num_classes = paths[name]
24 
25     ds = ds_fn(p, image_set=image_set, transforms=transform)
26     return ds, num_classes
27 
28 
29 def get_transform(train):
30     base_size = 520
31     crop_size = 480
32 
33     min_size = int((0.5 if train else 1.0) * base_size)
34     max_size = int((2.0 if train else 1.0) * base_size)
35     transforms = []
36     transforms.append(T.RandomResize(min_size, max_size))
37     if train:
38         transforms.append(T.RandomHorizontalFlip(0.5))
39         transforms.append(T.RandomCrop(crop_size))
40     transforms.append(T.ToTensor())
41     transforms.append(T.Normalize(mean=[0.485, 0.456, 0.406],
42                                   std=[0.229, 0.224, 0.225]))
43 
44     return T.Compose(transforms)
45 
46 
47 def criterion(inputs, target):
48     losses = {}
49     for name, x in inputs.items():
50         losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)
51 
52     if len(losses) == 1:
53         return losses['out']
54 
55     return losses['out'] + 0.5 * losses['aux']
56 
57 
58 def evaluate(model, data_loader, device, num_classes):
59     model.eval()
60     confmat = utils.ConfusionMatrix(num_classes)
61     metric_logger = utils.MetricLogger(delimiter="  ")
62     header = 'Test:'
63     with torch.no_grad():
64         for image, target in metric_logger.log_every(data_loader, 100, header):
65             image, target = image.to(device), target.to(device)
66             output = model(image)
67             output = output['out']
68 
69             confmat.update(target.flatten(), output.argmax(1).flatten())
70 
71         confmat.reduce_from_all_processes()
72 
73     return confmat
74 
75 
76 def train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, print_freq):
77     model.train()
78     metric_logger = utils.MetricLogger(delimiter="  ")
79     metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))
80     header = 'Epoch: [{}]'.format(epoch)
81     for image, target in metric_logger.log_every(data_loader, print_freq, header):
82         image, target = image.to(device), target.to(device)
83         output = model(image)
84         loss = criterion(output, target)
85 
86         optimizer.zero_grad()
87         loss.backward()
88         optimizer.step()
89 
90         lr_scheduler.step()
91 
92         metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
93 
94 
95 def main(args):
96     if args.output_dir:
97         utils.mkdir(args.output_dir)
98 
99     utils.init_distributed_mode(args)
100     print(args)
101 
102     device = torch.device(args.device)
103 
104     dataset, num_classes = get_dataset(args.dataset, "train", get_transform(train=True))
105     dataset_test, _ = get_dataset(args.dataset, "val", get_transform(train=False))
106 
107     if args.distributed:
108         train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
109         test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)
110     else:
111         train_sampler = torch.utils.data.RandomSampler(dataset)
112         test_sampler = torch.utils.data.SequentialSampler(dataset_test)
113 
114     data_loader = torch.utils.data.DataLoader(
115         dataset, batch_size=args.batch_size,
116         sampler=train_sampler, num_workers=args.workers,
117         collate_fn=utils.collate_fn, drop_last=True)
118 
119     data_loader_test = torch.utils.data.DataLoader(
120         dataset_test, batch_size=1,
121         sampler=test_sampler, num_workers=args.workers,
122         collate_fn=utils.collate_fn)
123 
124     model = torchvision.models.segmentation.__dict__[args.model](num_classes=num_classes,
125                                                                  aux_loss=args.aux_loss,
126                                                                  pretrained=args.pretrained)
127     model.to(device)
128     if args.distributed:
129         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
130 
131     if args.resume:
132         checkpoint = torch.load(args.resume, map_location='cpu')
133         model.load_state_dict(checkpoint['model'])
134 
135     model_without_ddp = model
136     if args.distributed:
137         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
138         model_without_ddp = model.module
139 
140     if args.test_only:
141         confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
142         print(confmat)
143         return
144 
145     params_to_optimize = [
146         {"params": [p for p in model_without_ddp.backbone.parameters() if p.requires_grad]},
147         {"params": [p for p in model_without_ddp.classifier.parameters() if p.requires_grad]},
148     ]
149     if args.aux_loss:
150         params = [p for p in model_without_ddp.aux_classifier.parameters() if p.requires_grad]
151         params_to_optimize.append({"params": params, "lr": args.lr * 10})
152     optimizer = torch.optim.SGD(
153         params_to_optimize,
154         lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
155 
156     lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
157         optimizer,
158         lambda x: (1 - x / (len(data_loader) * args.epochs)) ** 0.9)
159 
160     start_time = time.time()
161     for epoch in range(args.epochs):
162         if args.distributed:
163             train_sampler.set_epoch(epoch)
164         train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, args.print_freq)
165         confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)
166         print(confmat)
167         utils.save_on_master(
168             {
169                 'model': model_without_ddp.state_dict(),
170                 'optimizer': optimizer.state_dict(),
171                 'epoch': epoch,
172                 'args': args
173             },
174             os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
175 
176     total_time = time.time() - start_time
177     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
178     print('Training time {}'.format(total_time_str))
179 
180 
181 def parse_args():
182     import argparse
183     parser = argparse.ArgumentParser(description='PyTorch Segmentation Training')
184 
185     parser.add_argument('--dataset', default='voc', help='dataset')
186     parser.add_argument('--model', default='fcn_resnet101', help='model')
187     parser.add_argument('--aux-loss', action='store_true', help='auxiliar loss')
188     parser.add_argument('--device', default='cuda', help='device')
189     parser.add_argument('-b', '--batch-size', default=8, type=int)
190     parser.add_argument('--epochs', default=30, type=int, metavar='N',
191                         help='number of total epochs to run')
192 
193     parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
194                         help='number of data loading workers (default: 16)')
195     parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')
196     parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
197                         help='momentum')
198     parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
199                         metavar='W', help='weight decay (default: 1e-4)',
200                         dest='weight_decay')
201     parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
202     parser.add_argument('--output-dir', default='.', help='path where to save')
203     parser.add_argument('--resume', default='', help='resume from checkpoint')
204     parser.add_argument(
205         "--test-only",
206         dest="test_only",
207         help="Only test the model",
208         action="store_true",
209     )
210     parser.add_argument(
211         "--pretrained",
212         dest="pretrained",
213         help="Use pre-trained models from the modelzoo",
214         action="store_true",
215     )
216     # distributed training parameters
217     parser.add_argument('--world-size', default=1, type=int,
218                         help='number of distributed processes')
219     parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')
220 
221     args = parser.parse_args()
222     return args
223 
224 
225 if __name__ == "__main__":
226     args = parse_args()
227     main(args)
228
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
         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 @@
             {
                 'model': model_without_ddp.state_dict(),
                 'optimizer': optimizer.state_dict(),
+                'lr_scheduler': lr_scheduler.state_dict(),
                 'epoch': epoch,
                 'args': args
             },
@@ -201,6 +205,8 @@
     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",

{"golden_diff": "diff --git a/references/segmentation/train.py b/references/segmentation/train.py\n--- a/references/segmentation/train.py\n+++ b/references/segmentation/train.py\n@@ -128,10 +128,6 @@\n     if args.distributed:\n         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)\n \n-    if args.resume:\n-        checkpoint = torch.load(args.resume, map_location='cpu')\n-        model.load_state_dict(checkpoint['model'])\n-\n     model_without_ddp = model\n     if args.distributed:\n         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])\n@@ -157,8 +153,15 @@\n         optimizer,\n         lambda x: (1 - x / (len(data_loader) * args.epochs)) ** 0.9)\n \n+    if args.resume:\n+        checkpoint = torch.load(args.resume, map_location='cpu')\n+        model_without_ddp.load_state_dict(checkpoint['model'])\n+        optimizer.load_state_dict(checkpoint['optimizer'])\n+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n+        args.start_epoch = checkpoint['epoch'] + 1\n+\n     start_time = time.time()\n-    for epoch in range(args.epochs):\n+    for epoch in range(args.start_epoch, args.epochs):\n         if args.distributed:\n             train_sampler.set_epoch(epoch)\n         train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, args.print_freq)\n@@ -168,6 +171,7 @@\n             {\n                 'model': model_without_ddp.state_dict(),\n                 'optimizer': optimizer.state_dict(),\n+                'lr_scheduler': lr_scheduler.state_dict(),\n                 'epoch': epoch,\n                 'args': args\n             },\n@@ -201,6 +205,8 @@\n     parser.add_argument('--print-freq', default=10, type=int, help='print frequency')\n     parser.add_argument('--output-dir', default='.', help='path where to save')\n     parser.add_argument('--resume', default='', help='resume from checkpoint')\n+    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n+                        help='start epoch')\n     parser.add_argument(\n         \"--test-only\",\n         dest=\"test_only\",\n", "issue": "Resume in the Segmentation example\n## \ud83d\udc1b Bug\r\n\r\nThe 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:\r\n\r\n```\r\n    start_epoch = 0\r\n    if args.resume:\r\n        checkpoint = torch.load(args.resume, map_location='cpu')\r\n        model.load_state_dict(checkpoint['model'])\r\n        optimizer.load_state_dict(checkpoint['optimizer'])\r\n        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\r\n        start_epoch = checkpoint['epoch'] + 1\r\n```\r\n\r\nThen, [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:\r\n\r\n```\r\n    for epoch in range(start_epoch, args.epochs):\r\n```\r\n\r\nand\r\n\r\n```\r\nutils.save_on_master(\r\n            {\r\n                'model': model_without_ddp.state_dict(),\r\n                'optimizer': optimizer.state_dict(),\r\n                'lr_scheduler': lr_scheduler.state_dict(),\r\n                'epoch': epoch\r\n            },\r\n```\r\n\r\nBesides, 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`?\r\n\r\nThanks.\r\n\n", "before_files": [{"content": "import datetime\nimport os\nimport time\n\nimport torch\nimport torch.utils.data\nfrom torch import nn\nimport torchvision\n\nfrom coco_utils import get_coco\nimport transforms as T\nimport utils\n\n\ndef get_dataset(name, image_set, transform):\n    def sbd(*args, **kwargs):\n        return torchvision.datasets.SBDataset(*args, mode='segmentation', **kwargs)\n    paths = {\n        \"voc\": ('/datasets01/VOC/060817/', torchvision.datasets.VOCSegmentation, 21),\n        \"voc_aug\": ('/datasets01/SBDD/072318/', sbd, 21),\n        \"coco\": ('/datasets01/COCO/022719/', get_coco, 21)\n    }\n    p, ds_fn, num_classes = paths[name]\n\n    ds = ds_fn(p, image_set=image_set, transforms=transform)\n    return ds, num_classes\n\n\ndef get_transform(train):\n    base_size = 520\n    crop_size = 480\n\n    min_size = int((0.5 if train else 1.0) * base_size)\n    max_size = int((2.0 if train else 1.0) * base_size)\n    transforms = []\n    transforms.append(T.RandomResize(min_size, max_size))\n    if train:\n        transforms.append(T.RandomHorizontalFlip(0.5))\n        transforms.append(T.RandomCrop(crop_size))\n    transforms.append(T.ToTensor())\n    transforms.append(T.Normalize(mean=[0.485, 0.456, 0.406],\n                                  std=[0.229, 0.224, 0.225]))\n\n    return T.Compose(transforms)\n\n\ndef criterion(inputs, target):\n    losses = {}\n    for name, x in inputs.items():\n        losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)\n\n    if len(losses) == 1:\n        return losses['out']\n\n    return losses['out'] + 0.5 * losses['aux']\n\n\ndef evaluate(model, data_loader, device, num_classes):\n    model.eval()\n    confmat = utils.ConfusionMatrix(num_classes)\n    metric_logger = utils.MetricLogger(delimiter=\"  \")\n    header = 'Test:'\n    with torch.no_grad():\n        for image, target in metric_logger.log_every(data_loader, 100, header):\n            image, target = image.to(device), target.to(device)\n            output = model(image)\n            output = output['out']\n\n            confmat.update(target.flatten(), output.argmax(1).flatten())\n\n        confmat.reduce_from_all_processes()\n\n    return confmat\n\n\ndef train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, print_freq):\n    model.train()\n    metric_logger = utils.MetricLogger(delimiter=\"  \")\n    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))\n    header = 'Epoch: [{}]'.format(epoch)\n    for image, target in metric_logger.log_every(data_loader, print_freq, header):\n        image, target = image.to(device), target.to(device)\n        output = model(image)\n        loss = criterion(output, target)\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        lr_scheduler.step()\n\n        metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0][\"lr\"])\n\n\ndef main(args):\n    if args.output_dir:\n        utils.mkdir(args.output_dir)\n\n    utils.init_distributed_mode(args)\n    print(args)\n\n    device = torch.device(args.device)\n\n    dataset, num_classes = get_dataset(args.dataset, \"train\", get_transform(train=True))\n    dataset_test, _ = get_dataset(args.dataset, \"val\", get_transform(train=False))\n\n    if args.distributed:\n        train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)\n        test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)\n    else:\n        train_sampler = torch.utils.data.RandomSampler(dataset)\n        test_sampler = torch.utils.data.SequentialSampler(dataset_test)\n\n    data_loader = torch.utils.data.DataLoader(\n        dataset, batch_size=args.batch_size,\n        sampler=train_sampler, num_workers=args.workers,\n        collate_fn=utils.collate_fn, drop_last=True)\n\n    data_loader_test = torch.utils.data.DataLoader(\n        dataset_test, batch_size=1,\n        sampler=test_sampler, num_workers=args.workers,\n        collate_fn=utils.collate_fn)\n\n    model = torchvision.models.segmentation.__dict__[args.model](num_classes=num_classes,\n                                                                 aux_loss=args.aux_loss,\n                                                                 pretrained=args.pretrained)\n    model.to(device)\n    if args.distributed:\n        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)\n\n    if args.resume:\n        checkpoint = torch.load(args.resume, map_location='cpu')\n        model.load_state_dict(checkpoint['model'])\n\n    model_without_ddp = model\n    if args.distributed:\n        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])\n        model_without_ddp = model.module\n\n    if args.test_only:\n        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)\n        print(confmat)\n        return\n\n    params_to_optimize = [\n        {\"params\": [p for p in model_without_ddp.backbone.parameters() if p.requires_grad]},\n        {\"params\": [p for p in model_without_ddp.classifier.parameters() if p.requires_grad]},\n    ]\n    if args.aux_loss:\n        params = [p for p in model_without_ddp.aux_classifier.parameters() if p.requires_grad]\n        params_to_optimize.append({\"params\": params, \"lr\": args.lr * 10})\n    optimizer = torch.optim.SGD(\n        params_to_optimize,\n        lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n\n    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(\n        optimizer,\n        lambda x: (1 - x / (len(data_loader) * args.epochs)) ** 0.9)\n\n    start_time = time.time()\n    for epoch in range(args.epochs):\n        if args.distributed:\n            train_sampler.set_epoch(epoch)\n        train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, args.print_freq)\n        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)\n        print(confmat)\n        utils.save_on_master(\n            {\n                'model': model_without_ddp.state_dict(),\n                'optimizer': optimizer.state_dict(),\n                'epoch': epoch,\n                'args': args\n            },\n            os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))\n\n    total_time = time.time() - start_time\n    total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n    print('Training time {}'.format(total_time_str))\n\n\ndef parse_args():\n    import argparse\n    parser = argparse.ArgumentParser(description='PyTorch Segmentation Training')\n\n    parser.add_argument('--dataset', default='voc', help='dataset')\n    parser.add_argument('--model', default='fcn_resnet101', help='model')\n    parser.add_argument('--aux-loss', action='store_true', help='auxiliar loss')\n    parser.add_argument('--device', default='cuda', help='device')\n    parser.add_argument('-b', '--batch-size', default=8, type=int)\n    parser.add_argument('--epochs', default=30, type=int, metavar='N',\n                        help='number of total epochs to run')\n\n    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',\n                        help='number of data loading workers (default: 16)')\n    parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')\n    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n                        help='momentum')\n    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,\n                        metavar='W', help='weight decay (default: 1e-4)',\n                        dest='weight_decay')\n    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')\n    parser.add_argument('--output-dir', default='.', help='path where to save')\n    parser.add_argument('--resume', default='', help='resume from checkpoint')\n    parser.add_argument(\n        \"--test-only\",\n        dest=\"test_only\",\n        help=\"Only test the model\",\n        action=\"store_true\",\n    )\n    parser.add_argument(\n        \"--pretrained\",\n        dest=\"pretrained\",\n        help=\"Use pre-trained models from the modelzoo\",\n        action=\"store_true\",\n    )\n    # distributed training parameters\n    parser.add_argument('--world-size', default=1, type=int,\n                        help='number of distributed processes')\n    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')\n\n    args = parser.parse_args()\n    return args\n\n\nif __name__ == \"__main__\":\n    args = parse_args()\n    main(args)\n", "path": "references/segmentation/train.py"}], "after_files": [{"content": "import datetime\nimport os\nimport time\n\nimport torch\nimport torch.utils.data\nfrom torch import nn\nimport torchvision\n\nfrom coco_utils import get_coco\nimport transforms as T\nimport utils\n\n\ndef get_dataset(name, image_set, transform):\n    def sbd(*args, **kwargs):\n        return torchvision.datasets.SBDataset(*args, mode='segmentation', **kwargs)\n    paths = {\n        \"voc\": ('/datasets01/VOC/060817/', torchvision.datasets.VOCSegmentation, 21),\n        \"voc_aug\": ('/datasets01/SBDD/072318/', sbd, 21),\n        \"coco\": ('/datasets01/COCO/022719/', get_coco, 21)\n    }\n    p, ds_fn, num_classes = paths[name]\n\n    ds = ds_fn(p, image_set=image_set, transforms=transform)\n    return ds, num_classes\n\n\ndef get_transform(train):\n    base_size = 520\n    crop_size = 480\n\n    min_size = int((0.5 if train else 1.0) * base_size)\n    max_size = int((2.0 if train else 1.0) * base_size)\n    transforms = []\n    transforms.append(T.RandomResize(min_size, max_size))\n    if train:\n        transforms.append(T.RandomHorizontalFlip(0.5))\n        transforms.append(T.RandomCrop(crop_size))\n    transforms.append(T.ToTensor())\n    transforms.append(T.Normalize(mean=[0.485, 0.456, 0.406],\n                                  std=[0.229, 0.224, 0.225]))\n\n    return T.Compose(transforms)\n\n\ndef criterion(inputs, target):\n    losses = {}\n    for name, x in inputs.items():\n        losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)\n\n    if len(losses) == 1:\n        return losses['out']\n\n    return losses['out'] + 0.5 * losses['aux']\n\n\ndef evaluate(model, data_loader, device, num_classes):\n    model.eval()\n    confmat = utils.ConfusionMatrix(num_classes)\n    metric_logger = utils.MetricLogger(delimiter=\"  \")\n    header = 'Test:'\n    with torch.no_grad():\n        for image, target in metric_logger.log_every(data_loader, 100, header):\n            image, target = image.to(device), target.to(device)\n            output = model(image)\n            output = output['out']\n\n            confmat.update(target.flatten(), output.argmax(1).flatten())\n\n        confmat.reduce_from_all_processes()\n\n    return confmat\n\n\ndef train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, print_freq):\n    model.train()\n    metric_logger = utils.MetricLogger(delimiter=\"  \")\n    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))\n    header = 'Epoch: [{}]'.format(epoch)\n    for image, target in metric_logger.log_every(data_loader, print_freq, header):\n        image, target = image.to(device), target.to(device)\n        output = model(image)\n        loss = criterion(output, target)\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        lr_scheduler.step()\n\n        metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0][\"lr\"])\n\n\ndef main(args):\n    if args.output_dir:\n        utils.mkdir(args.output_dir)\n\n    utils.init_distributed_mode(args)\n    print(args)\n\n    device = torch.device(args.device)\n\n    dataset, num_classes = get_dataset(args.dataset, \"train\", get_transform(train=True))\n    dataset_test, _ = get_dataset(args.dataset, \"val\", get_transform(train=False))\n\n    if args.distributed:\n        train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)\n        test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)\n    else:\n        train_sampler = torch.utils.data.RandomSampler(dataset)\n        test_sampler = torch.utils.data.SequentialSampler(dataset_test)\n\n    data_loader = torch.utils.data.DataLoader(\n        dataset, batch_size=args.batch_size,\n        sampler=train_sampler, num_workers=args.workers,\n        collate_fn=utils.collate_fn, drop_last=True)\n\n    data_loader_test = torch.utils.data.DataLoader(\n        dataset_test, batch_size=1,\n        sampler=test_sampler, num_workers=args.workers,\n        collate_fn=utils.collate_fn)\n\n    model = torchvision.models.segmentation.__dict__[args.model](num_classes=num_classes,\n                                                                 aux_loss=args.aux_loss,\n                                                                 pretrained=args.pretrained)\n    model.to(device)\n    if args.distributed:\n        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)\n\n    model_without_ddp = model\n    if args.distributed:\n        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])\n        model_without_ddp = model.module\n\n    if args.test_only:\n        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)\n        print(confmat)\n        return\n\n    params_to_optimize = [\n        {\"params\": [p for p in model_without_ddp.backbone.parameters() if p.requires_grad]},\n        {\"params\": [p for p in model_without_ddp.classifier.parameters() if p.requires_grad]},\n    ]\n    if args.aux_loss:\n        params = [p for p in model_without_ddp.aux_classifier.parameters() if p.requires_grad]\n        params_to_optimize.append({\"params\": params, \"lr\": args.lr * 10})\n    optimizer = torch.optim.SGD(\n        params_to_optimize,\n        lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n\n    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(\n        optimizer,\n        lambda x: (1 - x / (len(data_loader) * args.epochs)) ** 0.9)\n\n    if args.resume:\n        checkpoint = torch.load(args.resume, map_location='cpu')\n        model_without_ddp.load_state_dict(checkpoint['model'])\n        optimizer.load_state_dict(checkpoint['optimizer'])\n        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n        args.start_epoch = checkpoint['epoch'] + 1\n\n    start_time = time.time()\n    for epoch in range(args.start_epoch, args.epochs):\n        if args.distributed:\n            train_sampler.set_epoch(epoch)\n        train_one_epoch(model, criterion, optimizer, data_loader, lr_scheduler, device, epoch, args.print_freq)\n        confmat = evaluate(model, data_loader_test, device=device, num_classes=num_classes)\n        print(confmat)\n        utils.save_on_master(\n            {\n                'model': model_without_ddp.state_dict(),\n                'optimizer': optimizer.state_dict(),\n                'lr_scheduler': lr_scheduler.state_dict(),\n                'epoch': epoch,\n                'args': args\n            },\n            os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))\n\n    total_time = time.time() - start_time\n    total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n    print('Training time {}'.format(total_time_str))\n\n\ndef parse_args():\n    import argparse\n    parser = argparse.ArgumentParser(description='PyTorch Segmentation Training')\n\n    parser.add_argument('--dataset', default='voc', help='dataset')\n    parser.add_argument('--model', default='fcn_resnet101', help='model')\n    parser.add_argument('--aux-loss', action='store_true', help='auxiliar loss')\n    parser.add_argument('--device', default='cuda', help='device')\n    parser.add_argument('-b', '--batch-size', default=8, type=int)\n    parser.add_argument('--epochs', default=30, type=int, metavar='N',\n                        help='number of total epochs to run')\n\n    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',\n                        help='number of data loading workers (default: 16)')\n    parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')\n    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n                        help='momentum')\n    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,\n                        metavar='W', help='weight decay (default: 1e-4)',\n                        dest='weight_decay')\n    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')\n    parser.add_argument('--output-dir', default='.', help='path where to save')\n    parser.add_argument('--resume', default='', help='resume from checkpoint')\n    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n                        help='start epoch')\n    parser.add_argument(\n        \"--test-only\",\n        dest=\"test_only\",\n        help=\"Only test the model\",\n        action=\"store_true\",\n    )\n    parser.add_argument(\n        \"--pretrained\",\n        dest=\"pretrained\",\n        help=\"Use pre-trained models from the modelzoo\",\n        action=\"store_true\",\n    )\n    # distributed training parameters\n    parser.add_argument('--world-size', default=1, type=int,\n                        help='number of distributed processes')\n    parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')\n\n    args = parser.parse_args()\n    return args\n\n\nif __name__ == \"__main__\":\n    args = parse_args()\n    main(args)\n", "path": "references/segmentation/train.py"}]}