We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
TypeError: can't pickle torch.distributed.ProcessGroupNCCL objects
Hi，
I meet a problem:

Traceback (most recent call last):
  File "train.py", line 394, in <module>
    train(hyp)
  File "train.py", line 331, in train
    torch.save(ckpt, last)
  File "/home/yy/anaconda3/lib/python3.6/site-packages/torch/serialization.py", line 328, in save
    _legacy_save(obj, opened_file, pickle_module, pickle_protocol)
  File "/home/yy/anaconda3/lib/python3.6/site-packages/torch/serialization.py", line 401, in _legacy_save
    pickler.dump(obj)
**TypeError: can't pickle torch.distributed.ProcessGroupNCCL objects**

Thanks!

environment:
ubuntu 16.04
GPU 2080Ti *4
pytorch 1.4.0


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `utils/torch_utils.py`
Content:
```
1 import math
2 import os
3 import time
4 from copy import deepcopy
5 
6 import torch
7 import torch.backends.cudnn as cudnn
8 import torch.nn as nn
9 import torch.nn.functional as F
10 import torchvision.models as models
11 
12 
13 def init_seeds(seed=0):
14     torch.manual_seed(seed)
15 
16     # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
17     if seed == 0:  # slower, more reproducible
18         cudnn.deterministic = True
19         cudnn.benchmark = False
20     else:  # faster, less reproducible
21         cudnn.deterministic = False
22         cudnn.benchmark = True
23 
24 
25 def select_device(device='', apex=False, batch_size=None):
26     # device = 'cpu' or '0' or '0,1,2,3'
27     cpu_request = device.lower() == 'cpu'
28     if device and not cpu_request:  # if device requested other than 'cpu'
29         os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
30         assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device  # check availablity
31 
32     cuda = False if cpu_request else torch.cuda.is_available()
33     if cuda:
34         c = 1024 ** 2  # bytes to MB
35         ng = torch.cuda.device_count()
36         if ng > 1 and batch_size:  # check that batch_size is compatible with device_count
37             assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)
38         x = [torch.cuda.get_device_properties(i) for i in range(ng)]
39         s = 'Using CUDA ' + ('Apex ' if apex else '')  # apex for mixed precision https://github.com/NVIDIA/apex
40         for i in range(0, ng):
41             if i == 1:
42                 s = ' ' * len(s)
43             print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
44                   (s, i, x[i].name, x[i].total_memory / c))
45     else:
46         print('Using CPU')
47 
48     print('')  # skip a line
49     return torch.device('cuda:0' if cuda else 'cpu')
50 
51 
52 def time_synchronized():
53     torch.cuda.synchronize() if torch.cuda.is_available() else None
54     return time.time()
55 
56 
57 def is_parallel(model):
58     # is model is parallel with DP or DDP
59     return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
60 
61 
62 def initialize_weights(model):
63     for m in model.modules():
64         t = type(m)
65         if t is nn.Conv2d:
66             pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
67         elif t is nn.BatchNorm2d:
68             m.eps = 1e-4
69             m.momentum = 0.03
70         elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
71             m.inplace = True
72 
73 
74 def find_modules(model, mclass=nn.Conv2d):
75     # finds layer indices matching module class 'mclass'
76     return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
77 
78 
79 def fuse_conv_and_bn(conv, bn):
80     # https://tehnokv.com/posts/fusing-batchnorm-and-conv/
81     with torch.no_grad():
82         # init
83         fusedconv = torch.nn.Conv2d(conv.in_channels,
84                                     conv.out_channels,
85                                     kernel_size=conv.kernel_size,
86                                     stride=conv.stride,
87                                     padding=conv.padding,
88                                     bias=True)
89 
90         # prepare filters
91         w_conv = conv.weight.clone().view(conv.out_channels, -1)
92         w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
93         fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
94 
95         # prepare spatial bias
96         if conv.bias is not None:
97             b_conv = conv.bias
98         else:
99             b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device)
100         b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
101         fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
102 
103         return fusedconv
104 
105 
106 def model_info(model, verbose=False):
107     # Plots a line-by-line description of a PyTorch model
108     n_p = sum(x.numel() for x in model.parameters())  # number parameters
109     n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients
110     if verbose:
111         print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
112         for i, (name, p) in enumerate(model.named_parameters()):
113             name = name.replace('module_list.', '')
114             print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
115                   (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
116 
117     try:  # FLOPS
118         from thop import profile
119         flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, 64, 64),), verbose=False)[0] / 1E9 * 2
120         fs = ', %.1f GFLOPS' % (flops * 100)  # 640x640 FLOPS
121     except:
122         fs = ''
123 
124     print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs))
125 
126 
127 def load_classifier(name='resnet101', n=2):
128     # Loads a pretrained model reshaped to n-class output
129     model = models.__dict__[name](pretrained=True)
130 
131     # Display model properties
132     input_size = [3, 224, 224]
133     input_space = 'RGB'
134     input_range = [0, 1]
135     mean = [0.485, 0.456, 0.406]
136     std = [0.229, 0.224, 0.225]
137     for x in [input_size, input_space, input_range, mean, std]:
138         print(x + ' =', eval(x))
139 
140     # Reshape output to n classes
141     filters = model.fc.weight.shape[1]
142     model.fc.bias = torch.nn.Parameter(torch.zeros(n), requires_grad=True)
143     model.fc.weight = torch.nn.Parameter(torch.zeros(n, filters), requires_grad=True)
144     model.fc.out_features = n
145     return model
146 
147 
148 def scale_img(img, ratio=1.0, same_shape=False):  # img(16,3,256,416), r=ratio
149     # scales img(bs,3,y,x) by ratio
150     h, w = img.shape[2:]
151     s = (int(h * ratio), int(w * ratio))  # new size
152     img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
153     if not same_shape:  # pad/crop img
154         gs = 32  # (pixels) grid size
155         h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
156     return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
157 
158 
159 class ModelEMA:
160     """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
161     Keep a moving average of everything in the model state_dict (parameters and buffers).
162     This is intended to allow functionality like
163     https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
164     A smoothed version of the weights is necessary for some training schemes to perform well.
165     E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use
166     RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA
167     smoothing of weights to match results. Pay attention to the decay constant you are using
168     relative to your update count per epoch.
169     To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but
170     disable validation of the EMA weights. Validation will have to be done manually in a separate
171     process, or after the training stops converging.
172     This class is sensitive where it is initialized in the sequence of model init,
173     GPU assignment and distributed training wrappers.
174     I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and single-GPU.
175     """
176 
177     def __init__(self, model, decay=0.9999, device=''):
178         # Create EMA
179         self.ema = deepcopy(model.module if is_parallel(model) else model)  # FP32 EMA
180         self.ema.eval()
181         self.updates = 0  # number of EMA updates
182         self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)
183         self.device = device  # perform ema on different device from model if set
184         if device:
185             self.ema.to(device)
186         for p in self.ema.parameters():
187             p.requires_grad_(False)
188 
189     def update(self, model):
190         # Update EMA parameters
191         with torch.no_grad():
192             self.updates += 1
193             d = self.decay(self.updates)
194 
195             msd = model.module.state_dict() if is_parallel(model) else model.state_dict()  # model state_dict
196             for k, v in self.ema.state_dict().items():
197                 if v.dtype.is_floating_point:
198                     v *= d
199                     v += (1. - d) * msd[k].detach()
200 
201     def update_attr(self, model):
202         # Update EMA attributes
203         for k, v in model.__dict__.items():
204             if not k.startswith('_') and k != 'module':
205                 setattr(self.ema, k, v)
206
```
--- 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/utils/torch_utils.py b/utils/torch_utils.py
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@@ -201,5 +201,5 @@
     def update_attr(self, model):
         # Update EMA attributes
         for k, v in model.__dict__.items():
-            if not k.startswith('_') and k != 'module':
+            if not k.startswith('_') and k not in ["process_group", "reducer"]:
                 setattr(self.ema, k, v)

{"golden_diff": "diff --git a/utils/torch_utils.py b/utils/torch_utils.py\n--- a/utils/torch_utils.py\n+++ b/utils/torch_utils.py\n@@ -201,5 +201,5 @@\n     def update_attr(self, model):\n         # Update EMA attributes\n         for k, v in model.__dict__.items():\n-            if not k.startswith('_') and k != 'module':\n+            if not k.startswith('_') and k not in [\"process_group\", \"reducer\"]:\n                 setattr(self.ema, k, v)\n", "issue": "TypeError: can't pickle torch.distributed.ProcessGroupNCCL objects\nHi\uff0c\r\nI meet a problem:\r\n\r\nTraceback (most recent call last):\r\n  File \"train.py\", line 394, in <module>\r\n    train(hyp)\r\n  File \"train.py\", line 331, in train\r\n    torch.save(ckpt, last)\r\n  File \"/home/yy/anaconda3/lib/python3.6/site-packages/torch/serialization.py\", line 328, in save\r\n    _legacy_save(obj, opened_file, pickle_module, pickle_protocol)\r\n  File \"/home/yy/anaconda3/lib/python3.6/site-packages/torch/serialization.py\", line 401, in _legacy_save\r\n    pickler.dump(obj)\r\n**TypeError: can't pickle torch.distributed.ProcessGroupNCCL objects**\r\n\r\nThanks!\r\n\r\nenvironment:\r\nubuntu 16.04\r\nGPU 2080Ti *4\r\npytorch 1.4.0\r\n\n", "before_files": [{"content": "import math\nimport os\nimport time\nfrom copy import deepcopy\n\nimport torch\nimport torch.backends.cudnn as cudnn\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.models as models\n\n\ndef init_seeds(seed=0):\n    torch.manual_seed(seed)\n\n    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html\n    if seed == 0:  # slower, more reproducible\n        cudnn.deterministic = True\n        cudnn.benchmark = False\n    else:  # faster, less reproducible\n        cudnn.deterministic = False\n        cudnn.benchmark = True\n\n\ndef select_device(device='', apex=False, batch_size=None):\n    # device = 'cpu' or '0' or '0,1,2,3'\n    cpu_request = device.lower() == 'cpu'\n    if device and not cpu_request:  # if device requested other than 'cpu'\n        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable\n        assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device  # check availablity\n\n    cuda = False if cpu_request else torch.cuda.is_available()\n    if cuda:\n        c = 1024 ** 2  # bytes to MB\n        ng = torch.cuda.device_count()\n        if ng > 1 and batch_size:  # check that batch_size is compatible with device_count\n            assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)\n        x = [torch.cuda.get_device_properties(i) for i in range(ng)]\n        s = 'Using CUDA ' + ('Apex ' if apex else '')  # apex for mixed precision https://github.com/NVIDIA/apex\n        for i in range(0, ng):\n            if i == 1:\n                s = ' ' * len(s)\n            print(\"%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)\" %\n                  (s, i, x[i].name, x[i].total_memory / c))\n    else:\n        print('Using CPU')\n\n    print('')  # skip a line\n    return torch.device('cuda:0' if cuda else 'cpu')\n\n\ndef time_synchronized():\n    torch.cuda.synchronize() if torch.cuda.is_available() else None\n    return time.time()\n\n\ndef is_parallel(model):\n    # is model is parallel with DP or DDP\n    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)\n\n\ndef initialize_weights(model):\n    for m in model.modules():\n        t = type(m)\n        if t is nn.Conv2d:\n            pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n        elif t is nn.BatchNorm2d:\n            m.eps = 1e-4\n            m.momentum = 0.03\n        elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:\n            m.inplace = True\n\n\ndef find_modules(model, mclass=nn.Conv2d):\n    # finds layer indices matching module class 'mclass'\n    return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]\n\n\ndef fuse_conv_and_bn(conv, bn):\n    # https://tehnokv.com/posts/fusing-batchnorm-and-conv/\n    with torch.no_grad():\n        # init\n        fusedconv = torch.nn.Conv2d(conv.in_channels,\n                                    conv.out_channels,\n                                    kernel_size=conv.kernel_size,\n                                    stride=conv.stride,\n                                    padding=conv.padding,\n                                    bias=True)\n\n        # prepare filters\n        w_conv = conv.weight.clone().view(conv.out_channels, -1)\n        w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))\n        fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))\n\n        # prepare spatial bias\n        if conv.bias is not None:\n            b_conv = conv.bias\n        else:\n            b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device)\n        b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))\n        fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)\n\n        return fusedconv\n\n\ndef model_info(model, verbose=False):\n    # Plots a line-by-line description of a PyTorch model\n    n_p = sum(x.numel() for x in model.parameters())  # number parameters\n    n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients\n    if verbose:\n        print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))\n        for i, (name, p) in enumerate(model.named_parameters()):\n            name = name.replace('module_list.', '')\n            print('%5g %40s %9s %12g %20s %10.3g %10.3g' %\n                  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))\n\n    try:  # FLOPS\n        from thop import profile\n        flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, 64, 64),), verbose=False)[0] / 1E9 * 2\n        fs = ', %.1f GFLOPS' % (flops * 100)  # 640x640 FLOPS\n    except:\n        fs = ''\n\n    print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs))\n\n\ndef load_classifier(name='resnet101', n=2):\n    # Loads a pretrained model reshaped to n-class output\n    model = models.__dict__[name](pretrained=True)\n\n    # Display model properties\n    input_size = [3, 224, 224]\n    input_space = 'RGB'\n    input_range = [0, 1]\n    mean = [0.485, 0.456, 0.406]\n    std = [0.229, 0.224, 0.225]\n    for x in [input_size, input_space, input_range, mean, std]:\n        print(x + ' =', eval(x))\n\n    # Reshape output to n classes\n    filters = model.fc.weight.shape[1]\n    model.fc.bias = torch.nn.Parameter(torch.zeros(n), requires_grad=True)\n    model.fc.weight = torch.nn.Parameter(torch.zeros(n, filters), requires_grad=True)\n    model.fc.out_features = n\n    return model\n\n\ndef scale_img(img, ratio=1.0, same_shape=False):  # img(16,3,256,416), r=ratio\n    # scales img(bs,3,y,x) by ratio\n    h, w = img.shape[2:]\n    s = (int(h * ratio), int(w * ratio))  # new size\n    img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize\n    if not same_shape:  # pad/crop img\n        gs = 32  # (pixels) grid size\n        h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]\n    return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean\n\n\nclass ModelEMA:\n    \"\"\" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models\n    Keep a moving average of everything in the model state_dict (parameters and buffers).\n    This is intended to allow functionality like\n    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage\n    A smoothed version of the weights is necessary for some training schemes to perform well.\n    E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use\n    RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA\n    smoothing of weights to match results. Pay attention to the decay constant you are using\n    relative to your update count per epoch.\n    To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but\n    disable validation of the EMA weights. Validation will have to be done manually in a separate\n    process, or after the training stops converging.\n    This class is sensitive where it is initialized in the sequence of model init,\n    GPU assignment and distributed training wrappers.\n    I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and single-GPU.\n    \"\"\"\n\n    def __init__(self, model, decay=0.9999, device=''):\n        # Create EMA\n        self.ema = deepcopy(model.module if is_parallel(model) else model)  # FP32 EMA\n        self.ema.eval()\n        self.updates = 0  # number of EMA updates\n        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)\n        self.device = device  # perform ema on different device from model if set\n        if device:\n            self.ema.to(device)\n        for p in self.ema.parameters():\n            p.requires_grad_(False)\n\n    def update(self, model):\n        # Update EMA parameters\n        with torch.no_grad():\n            self.updates += 1\n            d = self.decay(self.updates)\n\n            msd = model.module.state_dict() if is_parallel(model) else model.state_dict()  # model state_dict\n            for k, v in self.ema.state_dict().items():\n                if v.dtype.is_floating_point:\n                    v *= d\n                    v += (1. - d) * msd[k].detach()\n\n    def update_attr(self, model):\n        # Update EMA attributes\n        for k, v in model.__dict__.items():\n            if not k.startswith('_') and k != 'module':\n                setattr(self.ema, k, v)\n", "path": "utils/torch_utils.py"}], "after_files": [{"content": "import math\nimport os\nimport time\nfrom copy import deepcopy\n\nimport torch\nimport torch.backends.cudnn as cudnn\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.models as models\n\n\ndef init_seeds(seed=0):\n    torch.manual_seed(seed)\n\n    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html\n    if seed == 0:  # slower, more reproducible\n        cudnn.deterministic = True\n        cudnn.benchmark = False\n    else:  # faster, less reproducible\n        cudnn.deterministic = False\n        cudnn.benchmark = True\n\n\ndef select_device(device='', apex=False, batch_size=None):\n    # device = 'cpu' or '0' or '0,1,2,3'\n    cpu_request = device.lower() == 'cpu'\n    if device and not cpu_request:  # if device requested other than 'cpu'\n        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable\n        assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device  # check availablity\n\n    cuda = False if cpu_request else torch.cuda.is_available()\n    if cuda:\n        c = 1024 ** 2  # bytes to MB\n        ng = torch.cuda.device_count()\n        if ng > 1 and batch_size:  # check that batch_size is compatible with device_count\n            assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)\n        x = [torch.cuda.get_device_properties(i) for i in range(ng)]\n        s = 'Using CUDA ' + ('Apex ' if apex else '')  # apex for mixed precision https://github.com/NVIDIA/apex\n        for i in range(0, ng):\n            if i == 1:\n                s = ' ' * len(s)\n            print(\"%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)\" %\n                  (s, i, x[i].name, x[i].total_memory / c))\n    else:\n        print('Using CPU')\n\n    print('')  # skip a line\n    return torch.device('cuda:0' if cuda else 'cpu')\n\n\ndef time_synchronized():\n    torch.cuda.synchronize() if torch.cuda.is_available() else None\n    return time.time()\n\n\ndef is_parallel(model):\n    # is model is parallel with DP or DDP\n    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)\n\n\ndef initialize_weights(model):\n    for m in model.modules():\n        t = type(m)\n        if t is nn.Conv2d:\n            pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n        elif t is nn.BatchNorm2d:\n            m.eps = 1e-4\n            m.momentum = 0.03\n        elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:\n            m.inplace = True\n\n\ndef find_modules(model, mclass=nn.Conv2d):\n    # finds layer indices matching module class 'mclass'\n    return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]\n\n\ndef fuse_conv_and_bn(conv, bn):\n    # https://tehnokv.com/posts/fusing-batchnorm-and-conv/\n    with torch.no_grad():\n        # init\n        fusedconv = torch.nn.Conv2d(conv.in_channels,\n                                    conv.out_channels,\n                                    kernel_size=conv.kernel_size,\n                                    stride=conv.stride,\n                                    padding=conv.padding,\n                                    bias=True)\n\n        # prepare filters\n        w_conv = conv.weight.clone().view(conv.out_channels, -1)\n        w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))\n        fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))\n\n        # prepare spatial bias\n        if conv.bias is not None:\n            b_conv = conv.bias\n        else:\n            b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device)\n        b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))\n        fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)\n\n        return fusedconv\n\n\ndef model_info(model, verbose=False):\n    # Plots a line-by-line description of a PyTorch model\n    n_p = sum(x.numel() for x in model.parameters())  # number parameters\n    n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients\n    if verbose:\n        print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))\n        for i, (name, p) in enumerate(model.named_parameters()):\n            name = name.replace('module_list.', '')\n            print('%5g %40s %9s %12g %20s %10.3g %10.3g' %\n                  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))\n\n    try:  # FLOPS\n        from thop import profile\n        flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, 64, 64),), verbose=False)[0] / 1E9 * 2\n        fs = ', %.1f GFLOPS' % (flops * 100)  # 640x640 FLOPS\n    except:\n        fs = ''\n\n    print('Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs))\n\n\ndef load_classifier(name='resnet101', n=2):\n    # Loads a pretrained model reshaped to n-class output\n    model = models.__dict__[name](pretrained=True)\n\n    # Display model properties\n    input_size = [3, 224, 224]\n    input_space = 'RGB'\n    input_range = [0, 1]\n    mean = [0.485, 0.456, 0.406]\n    std = [0.229, 0.224, 0.225]\n    for x in [input_size, input_space, input_range, mean, std]:\n        print(x + ' =', eval(x))\n\n    # Reshape output to n classes\n    filters = model.fc.weight.shape[1]\n    model.fc.bias = torch.nn.Parameter(torch.zeros(n), requires_grad=True)\n    model.fc.weight = torch.nn.Parameter(torch.zeros(n, filters), requires_grad=True)\n    model.fc.out_features = n\n    return model\n\n\ndef scale_img(img, ratio=1.0, same_shape=False):  # img(16,3,256,416), r=ratio\n    # scales img(bs,3,y,x) by ratio\n    h, w = img.shape[2:]\n    s = (int(h * ratio), int(w * ratio))  # new size\n    img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize\n    if not same_shape:  # pad/crop img\n        gs = 32  # (pixels) grid size\n        h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]\n    return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean\n\n\nclass ModelEMA:\n    \"\"\" Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models\n    Keep a moving average of everything in the model state_dict (parameters and buffers).\n    This is intended to allow functionality like\n    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage\n    A smoothed version of the weights is necessary for some training schemes to perform well.\n    E.g. Google's hyper-params for training MNASNet, MobileNet-V3, EfficientNet, etc that use\n    RMSprop with a short 2.4-3 epoch decay period and slow LR decay rate of .96-.99 requires EMA\n    smoothing of weights to match results. Pay attention to the decay constant you are using\n    relative to your update count per epoch.\n    To keep EMA from using GPU resources, set device='cpu'. This will save a bit of memory but\n    disable validation of the EMA weights. Validation will have to be done manually in a separate\n    process, or after the training stops converging.\n    This class is sensitive where it is initialized in the sequence of model init,\n    GPU assignment and distributed training wrappers.\n    I've tested with the sequence in my own train.py for torch.DataParallel, apex.DDP, and single-GPU.\n    \"\"\"\n\n    def __init__(self, model, decay=0.9999, device=''):\n        # Create EMA\n        self.ema = deepcopy(model.module if is_parallel(model) else model)  # FP32 EMA\n        self.ema.eval()\n        self.updates = 0  # number of EMA updates\n        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))  # decay exponential ramp (to help early epochs)\n        self.device = device  # perform ema on different device from model if set\n        if device:\n            self.ema.to(device)\n        for p in self.ema.parameters():\n            p.requires_grad_(False)\n\n    def update(self, model):\n        # Update EMA parameters\n        with torch.no_grad():\n            self.updates += 1\n            d = self.decay(self.updates)\n\n            msd = model.module.state_dict() if is_parallel(model) else model.state_dict()  # model state_dict\n            for k, v in self.ema.state_dict().items():\n                if v.dtype.is_floating_point:\n                    v *= d\n                    v += (1. - d) * msd[k].detach()\n\n    def update_attr(self, model):\n        # Update EMA attributes\n        for k, v in model.__dict__.items():\n            if not k.startswith('_') and k not in [\"process_group\", \"reducer\"]:\n                setattr(self.ema, k, v)\n", "path": "utils/torch_utils.py"}]}