PyTorch 1.7.0 Compatibility Updates

hubconf.py

@@ -108,3 +108,11 @@ def yolov5x(pretrained=False, channels=3, classes=80):
 
 if __name__ == '__main__':
     model = create(name='yolov5s', pretrained=True, channels=3, classes=80)  # example
+    model = model.fuse().eval().autoshape()  # for autoshaping of PIL/cv2/np inputs and NMS
+
+    # Verify inference
+    from PIL import Image
+
+    img = Image.open('inference/images/zidane.jpg')
+    y = model(img)
+    print(y[0].shape)

models/experimental.py

@@ -136,6 +136,13 @@ def attempt_load(weights, map_location=None):
         attempt_download(w)
         model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval())  # load FP32 model
 
+    # Compatibility updates
+    for m in model.modules():
+        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
+            m.inplace = True  # pytorch 1.7.0 compatibility
+        elif type(m) is Conv:
+            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+
     if len(model) == 1:
         return model[-1]  # return model
     else:

models/yolo.py

@@ -165,7 +165,6 @@ def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
         print('Fusing layers... ')
         for m in self.model.modules():
             if type(m) is Conv and hasattr(m, 'bn'):
-                m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatability
                 m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                 delattr(m, 'bn')  # remove batchnorm
                 m.forward = m.fuseforward  # update forward

utils/torch_utils.py

@@ -74,7 +74,7 @@ def initialize_weights(model):
         elif t is nn.BatchNorm2d:
             m.eps = 1e-3
             m.momentum = 0.03
-        elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
+        elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
             m.inplace = True