mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Rearrange export input after checks (#1118)

Browse Source 
img size checks are warnings rather than errors, so current implementation allows improperly formed model inputs.
This commit is contained in:
Glenn Jocher 2020-10-11 16:23:36 +02:00 committed by GitHub
parent d87cf7ecb9
commit d45e349a1e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
models/export.py
View File

@ -29,9 +29,6 @@ if __name__ == '__main__':
set_logging()
t = time.time()
# Input
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection
# Load PyTorch model
model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model
labels = model.names
@ -40,6 +37,9 @@ if __name__ == '__main__':
gs = int(max(model.stride)) # grid size (max stride)
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
# Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size) # image size(1,3,320,192) iDetection
# Update model
for k, m in model.named_modules():
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility