Pytorch 深度强化学习模型训练速度慢

最近一直在用Pytorch来训练深度强化学习模型，但是速度一直很慢，Gpu利用率也很低。

一、起初开始在训练参数 batch_size = 200, graph_size = 40, epoch_size=100000, 训练速度一个epoch要4h，人麻了…

二、经过网上查找资料后，对batch_size 和训练数据DataLoader 的pin_memory，num_workers参数进行了修改调试，结果还是不行，虽然GPU利用率有提高一点，但是速度还是没改进。
三、之后才发现是强化学习的问题环境代码有问题。

问题1.在问题环境中用了两个 双重for循环来对解决数据，源代码如下：

for i in range(arrive_time.shape[0]):
      if arrive_time[i] > tw[rng, a, 1][i]:
          delay_time[i] = arrive_time[i] - tw[rng, a, 1][i]
      else:
         delay_time[i] = 0

主要是 利用for循环来判断符合一些条件的值并进行处理，虽然类型都是Tensor，调试时也没感觉耗费很多时间，但是把这部分代码去掉后，该做Tensor的函数进行操作后模型训练时间有了提高。
修改后的代码：利用where()函数来进行判断并处理

delay_time_temp = arrive_time - tw[rng, a, 1]
zero2 = torch.zeros_like(delay_time_temp)
delay_time = torch.where(delay_time_temp < 0, zero2, delay_time_temp)

问题2. 有冗余的无意义代码存在
在调试代码的过程中自己有写很多没用的tensor变量和其他一些代码，之后忘记注释掉。这部分也浪费时间。在注释之后时间又有了一部分提高。在相同参数下一个epoch 时间到了十几分钟。

发愁了好几天，记录一下，在训练模型时要避免使用过多的for 语句或者一些冗余的代码，对于tensor 还是要用相关函数来进行处理。具体原因还不太懂，但是速度个人还是不满意，提高graph_size 后速度还是会到30分钟左右，希望还能在优化。