diff --git a/hubconf.py b/hubconf.py index cd14863ca8b41eb5570eb00191ac3b57a58a7838..cc210528c087913e804c995c3a99edc8f8268c61 100644 --- a/hubconf.py +++ b/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) diff --git a/models/experimental.py b/models/experimental.py index 0b61027b9d2f9aac2d3c88a666a5b725d976d4ce..a2908a15cf320d08b059f4d44e2cff864ee7a3e2 100644 --- a/models/experimental.py +++ b/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: diff --git a/models/yolo.py b/models/yolo.py index 0d46054ed21c65828aca1a5b0aad0fdfbfcf1b00..e1c30baa271dc68610c6539248dac8b8025e7cdf 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -165,7 +165,6 @@ class Model(nn.Module): 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 diff --git a/utils/torch_utils.py b/utils/torch_utils.py index f6818238452f77f95911fd5f31d12dff95df38a7..25eff07f3f444d9083682761d7adfe90e8f3d026 100644 --- a/utils/torch_utils.py +++ b/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