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

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
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -8,7 +8,7 @@ from copy import deepcopy
import numpy as np
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
@ -29,7 +29,7 @@ def check_train_batch_size(model, imgsz=640, amp=True):
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.