diff --git a/models/yolo.py b/models/yolo.py index 0fa2db91e82b..80ff83e16085 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -20,7 +20,7 @@ from models.common import * from models.experimental import * from utils.autoanchor import check_anchor_order -from utils.general import check_yaml, make_divisible, print_args, set_logging +from utils.general import check_yaml, make_divisible, print_args, set_logging, check_version from utils.plots import feature_visualization from utils.torch_utils import copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, \ select_device, time_sync @@ -74,7 +74,10 @@ def forward(self, x): def _make_grid(self, nx=20, ny=20, i=0): d = self.anchors[i].device - yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)]) + if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility + yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)], indexing='ij') + else: + yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)]) grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float() anchor_grid = (self.anchors[i].clone() * self.stride[i]) \ .view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float() diff --git a/utils/augmentations.py b/utils/augmentations.py index 04192d1ec5cd..b3cbbf913b65 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -20,7 +20,7 @@ def __init__(self): self.transform = None try: import albumentations as A - check_version(A.__version__, '1.0.3') # version requirement + check_version(A.__version__, '1.0.3', hard=True) # version requirement self.transform = A.Compose([ A.Blur(p=0.01), diff --git a/utils/general.py b/utils/general.py index f22908907fd0..667af63e4044 100755 --- a/utils/general.py +++ b/utils/general.py @@ -220,14 +220,17 @@ def check_git_status(): def check_python(minimum='3.6.2'): # Check current python version vs. required python version - check_version(platform.python_version(), minimum, name='Python ') + check_version(platform.python_version(), minimum, name='Python ', hard=True) -def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False): +def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False): # Check version vs. required version current, minimum = (pkg.parse_version(x) for x in (current, minimum)) - result = (current == minimum) if pinned else (current >= minimum) - assert result, f'{name}{minimum} required by YOLOv5, but {name}{current} is currently installed' + result = (current == minimum) if pinned else (current >= minimum) # bool + if hard: # assert min requirements met + assert result, f'{name}{minimum} required by YOLOv5, but {name}{current} is currently installed' + else: + return result @try_except