Default AutoBatch 0.8 fraction (#9212)
This commit is contained in:
1 parent
3c64d89104
commit
91a81d48fa
2 files changed
+2
-2
No files matched your search
+1
-1
@@ -47,7 +47,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
|
|||||||
model = DetectMultiBackend(path, device=device, fuse=autoshape) # detection model
|
model = DetectMultiBackend(path, device=device, fuse=autoshape) # detection model
|
||||||
if autoshape:
|
if autoshape:
|
||||||
if model.pt and isinstance(model.model, ClassificationModel):
|
if model.pt and isinstance(model.model, ClassificationModel):
|
||||||
LOGGER.warning('WARNING: YOLOv5 v6.2 ClassificationModel is not yet AutoShape compatible. '
|
LOGGER.warning('WARNING: ⚠️ YOLOv5 v6.2 ClassificationModel is not yet AutoShape compatible. '
|
||||||
'You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224).')
|
'You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224).')
|
||||||
else:
|
else:
|
||||||
model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS
|
model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS
|
||||||
|
|||||||
+1
-1
@@ -18,7 +18,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.9, batch_size=16):
|
def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
|
||||||
# Automatically estimate best batch size to use `fraction` of available CUDA memory
|
# Automatically estimate best batch size to use `fraction` of available CUDA memory
|
||||||
# Usage:
|
# Usage:
|
||||||
# import torch
|
# import torch
|
||||||
|
|||||||
Reference in new issue
Block a user