xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix batch=-1 for exports (#3217)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-06-17 01:07:24 +02:00
committed by GitHub
parent d69a1e8046
commit 0dfafaf39b
3 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ultralytics/yolo/engine/model.py
View File

@ -335,7 +335,7 @@ class YOLO:
overrides['mode'] = 'export' overrides['mode'] = 'export'
if overrides.get('imgsz') is None: if overrides.get('imgsz') is None:
overrides['imgsz'] = self.model.args['imgsz'] # use trained imgsz unless custom value is passed overrides['imgsz'] = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
if overrides.get('batch') is None: if 'batch' not in kwargs:
overrides['batch'] = 1 # default to 1 if not modified overrides['batch'] = 1 # default to 1 if not modified
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides) args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.task = self.task args.task = self.task

2
ultralytics/yolo/engine/trainer.py
View File

@ -229,7 +229,7 @@ class BaseTrainer:
# Batch size # Batch size
if self.batch_size == -1: if self.batch_size == -1:
if RANK == -1: # single-GPU only, estimate best batch size if RANK == -1: # single-GPU only, estimate best batch size
self.batch_size = check_train_batch_size(self.model, self.args.imgsz, self.amp) self.args.batch = self.batch_size = check_train_batch_size(self.model, self.args.imgsz, self.amp)
else: else:
SyntaxError('batch=-1 to use AutoBatch is only available in Single-GPU training. ' SyntaxError('batch=-1 to use AutoBatch is only available in Single-GPU training. '
'Please pass a valid batch size value for Multi-GPU DDP training, i.e. batch=16') 'Please pass a valid batch size value for Multi-GPU DDP training, i.e. batch=16')

4
ultralytics/yolo/utils/autobatch.py
View File

@ -8,7 +8,7 @@ from copy import deepcopy
import numpy as np import numpy as np
import torch import torch
from ultralytics.yolo.utils import LOGGER, colorstr from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, colorstr
from ultralytics.yolo.utils.torch_utils import profile from ultralytics.yolo.utils.torch_utils import profile
@ -29,7 +29,7 @@ def check_train_batch_size(model, imgsz=640, amp=True):
return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size
def autobatch(model, imgsz=640, fraction=0.67, batch_size=16): def autobatch(model, imgsz=640, fraction=0.67, batch_size=DEFAULT_CFG.batch):
""" """
Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory. Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.