I set the model hidden layer a litter bigger(like 256), and use load_state_dict
in the subprocess.
When I give a multi-dimensional tensor to the model, the process will be terminated without any exception.
I write a demo to reproduce the error.
import torch.multiprocessing as mp
import torch
import torch.nn as nn
class AC(nn.Module):
def __init__(self, features_n, actions_n):
super(AC, self).__init__()
# if the hidden layer cells are lower, for example, 128, no error occurs.
self.hidden_layer_cells = 256
self.l1 = nn.Linear(features_n, self.hidden_layer_cells)
self.l2 = nn.Linear(self.hidden_layer_cells, self.hidden_layer_cells)
self.actor_linear = nn.Linear(self.hidden_layer_cells, actions_n)
self.critic_linear = nn.Linear(self.hidden_layer_cells, 1)
def forward(self, inputs):
x = torch.tanh(self.l1(inputs))
x = torch.tanh(self.l2(x))
pi = self.actor_linear(x)
q = self.critic_linear(x)
return pi, q
class Worker2(mp.Process):
def __init__(self) -> None:
super(Worker2, self).__init__()
self.t = AC(10, 10)
self.tt = AC(10, 10)
# if I load state dict from exist model, it will be terminated when passing Multidimensional tensor
self.tt.load_state_dict(self.t.state_dict())
def run(self):
while True:
s = torch.ones(size=(1, 10))
a = self.t(s)
ss = torch.cat((s, s))
# this line will terminate the process
aa = self.t(ss)
w = Worker2()
w.start()
w.join()