HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Simplify argument names (#141)

Browse source
This commit is contained in:
Glenn Jocher 2023-01-03 19:47:56 +01:00 • committed by GitHub
parent d74de2582c
commit 3c18aa5278
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 22 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

6
ultralytics/yolo/engine/exporter.py
View file

@ -164,14 +164,14 @@ class Exporter:
assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic'
# Checks
# if self.args.batch_size == model.args['batch_size']: # user has not modified training batch_size
self.args.batch_size = 1
# if self.args.batch == model.args['batch_size']: # user has not modified training batch_size
self.args.batch = 1
self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size
if self.args.optimize:
assert self.device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
# Input
im = torch.zeros(self.args.batch_size, 3, *self.imgsz).to(self.device)
im = torch.zeros(self.args.batch, 3, *self.imgsz).to(self.device)
file = Path(getattr(model, 'pt_path', None) or model.yaml['yaml_file'])
if file.suffix == '.yaml':
file = Path(file.name)