num_workers(dataloader)、batch_size、gpu gpu利用率低解决方法
dataset = SimpleDataset() dataloader = DataLoader(dataset, num_workers=2, batch_size=3) for batch in dataloader: print(batch)
主进程初始化
dataloader = DataLoader(dataset, num_workers=2, batch_size=3)
- 创建
num_workers
个不同的子进程。 - 为每个子进程维护一个
index_queue
,用于发送需要采集数据的index。注意,队列每个元素都是一个batch的index的数组。 - 所有子进程共享一个
data_queue
,用于存放子进程返回的数据。 - 为每个子进程的
index_queue
依次存放两个batch的数据。比如本代码,会为0号worker存放[(0,1,2),(6,7,8)]
,为1号worker存放[(3,4,5),(9,10,11)]
。依次存放是为了尽量保证取得数据的顺序和采样器生成的顺序是一致的。而存放两组数据是为了实现数据预取的功能(即取第i个batch的数据时,i+1个batch的数据已经读好,并开始缓存i+2个batch的数据。
基于上述分析,我们可以很容易得到以下几个工程经验:
- 一个完整的batch由一个worker产生,而不是由多个worker共同产生。这样做主要是为了尽可能减少进程通信的次数。
- 内存占用有多大?正常情况下,内存中稳定存储
2*num_workers
个batch的数据。 但由于乱序缓存机制的存在,实际会比这个大一些。最坏情况是某个worker卡住始终不返回数据,那就会一直缓存后续数据,一直读取,直到爆内存。个人认为这个设计并不好,为何不每个worker维护自己的data_queue? - 数据预取的trick还需要加吗?之前在很多文章看到自己定义DataLoader,然后手动prefetch下一个batch的数据以加速。个人认为由于双缓存机制,这项trick已经没用了。
- 在设计
Dataset
时,应该让每次__getitem__
的时间尽可能均匀,这样不会浪费过多的内存在乱序缓存机制上,可以有效降低内存占用。
问题:
在跑tsm模型的时候,发现gpu的利用率低,而且中间有间断,num_workers设置的是16,batch_size16,有8个显卡(后来检查才发现有56个cpu)
然后把num_workers改成56,batchsize也改为56,解决!!!gpu利用率一下子提升了
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 阿里最新开源QwQ-32B,效果媲美deepseek-r1满血版,部署成本又又又降低了!
· SQL Server 2025 AI相关能力初探
· 单线程的Redis速度为什么快?
· AI编程工具终极对决:字节Trae VS Cursor,谁才是开发者新宠?
· 开源Multi-agent AI智能体框架aevatar.ai,欢迎大家贡献代码
2020-03-25 手写数字识别
2020-03-25 一个讲pyYaml的博客
2020-03-25 tf.cast()
2020-03-25 argparse基本用法