直接在gpu上定义tensor而减少cpu到gpu的copy!!!
问题
为了实现复杂的网络结构和损失函数,自定义forward和loss的计算,发现cpu占用很高,挂了几个训练就占满了所有cpu。
把运算都用.to(device)
放到gpu上运算,发现还是cpu占用高,并且有更慢的趋势。用性能分析工具分析发现是aten::_to_copy
aten::copy_
cpu占用时间率高。
分析
为了验证是使用了.to(device)
导致的,编写了程序,对比.to('cuda')
和直接定义torch.tensor([],device='cuda')
,发现torch.tensor([],device='cuda')
时gpu利用率接近100%,而.to('cuda')
gpu利用率还不到20%。
结论
直接在gpu上定义变量torch.tensor([],device='cuda')