「解析」Pytorch 自动计算 batchsize

日志是一个十分必要的操作，有助于后期分析实验结果，特别是在多台不同环境下训练，为了区分，还是十分有必要记录相关平台信息的，比如 hostname，Python版本信息，Pytorch版本信息等！

import socket
import platform'Python-{platform.python_version()}'				# Python 版本
'torch-{torch.__version__}'							# Pytorch 版本
socket.gethostname()								# 主机名s = f'YOLOv5 🚀 {time_date()} Python-{platform.python_version()} torch-{torch.__version__} '

1、查看显卡信息

首先需要获取显卡信息，查看显存 File path： yolov5-7.0/utils/autobatch.py

torch.cuda.get_device_properties(0)			# 获取显卡属性
>>> _CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24563MB, multi_processor_count=128)
'name:' 显卡型号
'major:'
'minor:'
'total_memory:' 总显存
'multi_processor_count: 'torch.cuda.empty_cache()"""-----------------------------------------------------------------------------------------"""
properties = torch.cuda.get_device_properties(device)		# 获取cuda:i属性a = torch.cuda.memory_allocated(device) / gb  				# GiB allocated	当前Tensor占用的显存 
# Returns the current GPU memory occupied by tensors in bytes for a given device.r = torch.cuda.memory_reserved(device) / gb  				# GiB reserved	总共占用的显存
# Returns the current GPU memory managed by the caching allocator in bytes for a given device.t = properties.total_memory / gb  							# GiB total
f = t - (r + a)  											# GiB freeprint(f'{cuda_i} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free')
>>> 'CUDA:0 (NVIDIA GeForce RTX 4090) 23.99G total, 0.09G reserved, 0.05G allocated, 23.85G free'

2、自动计算 batchsize

上面已经获取了显卡的占用显存，空余显存和总显存

首先创建个 batch_size 的列表，用于计算不同小batch时，显卡的使用情况，最后再通过 np.polyfit 拟合计算出最佳的 batch_size

# Profile batch sizes
batch_sizes = [1, 2, 4, 8, 16]
try:img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]results = profile(img, model, n=3, device=device)			# profile：yolov5-7.0/utils/torch_utils.py
except Exception as e:LOGGER.warning(f'{prefix}{e}')# Fit a solution
y = [x[2] for x in results if x]  					# memory [2]
p = np.polyfit(batch_sizes[:len(y)], y, deg=1)  	# first degree polynomial fit
b = int((f * fraction - p[1]) / p[0])  				# y intercept (optimal batch size)

• hasattr(): 函数用于判断对象是否包含对应的属性
• isinstance(): 用来判断一个函数是否是一个已知的类型，类似 type()
• torch.cuda.synchronize(): 等待当前设备上所有流中的所有核心完成
• torch.numel(input) → int 返回输入张量中元素的总数

通过 np.polyfit 拟合可以看出，显存占用情况基本随 batch_size 线性增加，其中 np.polyfit 的 deg参数表示待拟合多项式的次数，输出结果从最高次幂依次递减 。通过上述计算的 free 显存，乘以 设定比例，然后减去 偏置显存，然后除以 斜率既是对应的 batch_size，再对其取整获得即可

3、完整版代码