This commit is contained in:
dianjixz
2023-11-06 09:03:18 +08:00
parent a6f6ce424d
commit faa96f40bf
2 changed files with 2 additions and 2 deletions
+1 -1
View File
@@ -555,7 +555,7 @@ def run(
y = model(im) # dry runs y = model(im) # dry runs
if half and not coreml: if half and not coreml:
im, model = im.half(), model.half() # to FP16 im, model = im.half(), model.half() # to FP16
if os.environ['EXPORT_ARCH'] == 'ax620': if os.environ.get("EXPORT_ARCH", "local") == 'ax620':
metadata = {'stride': int(max(model.stride)), 'names': model.names} # model metadata metadata = {'stride': int(max(model.stride)), 'names': model.names} # model metadata
else: else:
shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape
+1 -1
View File
@@ -57,7 +57,7 @@ class Detect(nn.Module):
z = [] # inference output z = [] # inference output
for i in range(self.nl): for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv x[i] = self.m[i](x[i]) # conv
if 'EXPORT_ARCH' in os.environ and os.environ['EXPORT_ARCH'] == 'ax620': if os.environ.get("EXPORT_ARCH", "local") == 'ax620':
return x return x
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()