o Ďi @sTdZddlmZddlZddlZddlmZmZddl m Z ddd Z dd d Z dS)z Auto-batch utils )deepcopyN)LOGGERcolorstr)profileTcCsDtjj|tt||WdS1swYdS)N)torchcudaampautocast autobatchrtrain)modelimgszr rR/home/jeff/fluffinator/venv/lib/python3.10/site-packages/yolov5/utils/autobatch.pycheck_train_batch_sizes$r皙?csTtd}t|dt|j}|jdkr&t|d||Stjj j r7t|d||Sd}t | }tj |}|j|} tj ||} tj ||} | | | } t||d|jd| d d | d d | d d | d d gd} zfdd| D}t||d|d}Wnty}zt||WYd}~nd}~wwdd|D}tj| dt||dd}t| ||d|d}d|vr|d}|| |kr| t|dd}|dks|dkr|}t|dt||| | | }t|d|d|d| |d d| d d|ddd |S)!Nz AutoBatch: z)Computing optimal batch size for --imgsz cpuz0CUDA not detected, using default CPU batch-size uP ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size i@z (z) z.2fz G total, z G reserved, z G allocated, zG free)rcsg|] }t|dqS))rempty).0brrr 5szautobatch..r)ndevicecSsg|]}|r|dqS)rr)rxrrrr;sr)degriuVWARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.zUsing batch-size z for  zG/zG (dz.0fu%) ✅)rrinfonext parametersr typerbackendscudnn benchmarkstrupperrget_device_properties total_memorymemory_reservedmemory_allocatednamer Exceptionwarningnppolyfitlenintindexmaxpolyval)r rfraction batch_sizeprefixr gbd propertiestraf batch_sizesimgresultseyprirrrr sL       <  <r )rT)rrr) __doc__copyrnumpyr5ryolov5.utils.generalrryolov5.utils.torch_utilsrrr rrrrs