pytorch 自定义forward loss运算cpu占用高 在gpu上加速

直接在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')