I’ve implemented a simple DDQN network in pytorch and tensorflow. The network is quite shallow.
While the forward pass is much faster in PyTorch compared to TF, the back-propagation step is much slower compared to TF. Both backprop steps were done on the CPU.
Any ideas how to improve it.
The network part is:
def __init__(self, hidden_size_IP=100, hidden_size_rest=100, alpha=0.01, state_size=27, action_size=8, learning_rate=1e-6): super().__init__() # build hidden layers self.l1 = nn.Sequential(nn.Linear(in_features=500, out_features=400), nn.LeakyReLU(negative_slope=alpha)) self.l2 = nn.Sequential(nn.Linear(in_features=400, out_features=200), nn.LeakyReLU(negative_slope=alpha)) self.l3 = nn.Sequential(nn.Linear(in_features=200, out_features=200), nn.LeakyReLU(negative_slope=alpha)) # build output layer self.Qval = nn.Linear(in_features=200, out_features=24) def forward(self, observation): if isinstance(observation, np.ndarray): observation = torch.from_numpy(observation).float() out1 = self.l1(observation) out2 = self.l2(out1) out3 = self.l3(out2) qval = self.Qval(out3) return qval
and the backprop code can be, for example:
self.optimizer = optim.Adam(self.q_net.parameters(), lr=1e-4)
state_batch=torch.rand([64,500])
act_batch=np.randi(0,24,[64,1]
act_batch_torch=torch.as_tensor(act_batch)
label_batch = torch.rand([64,500])
Q=self.q_net.forward(state_batch).gather(1, act_batch_torch) # q_net is an instance of the network above
loss = mse_loss(input=Q, target=label_batch.detach())
loss.backward()
self.optimizer.zero_grad()
self.optimizer.step()
Note that since inference is much faster using the CPU, I’m also doing backprop on the CPU. I have tried transferring the network to the GPU, and then do a backprop on the GPU, but it turned out to be slower.
Any ideas why pyTorch is slower? How can I improve the speed for this type of shallow network?
I’m using pyTorch 1.0.