I'm trying to code an Asynchronous Actor Critic in PyTorch based on this repo: https://github.com/seungeunrho/minimalRL/blob/master/a3c.py but I'm changing the ActorCritic class to use the one I coded myself.
Basically I have a class A3C, an instance of it, global_model, with shared memory and I use torch.multiprocessing to open some Processes in order to train the model in parallel. In each process at the beginning I have to create a new instance of the model, called local_model, in order to proceed with the training, but the process gets stuck in the initialization of the local model even though the one of the global model works every time.
Trying to debugging it I can see that it enters the A3C.init function and the SharedActorCritic.init too, but there it stops just after I put the checkpoint print. However if I print whatever expression contains list(critic_param_gen) magically everything works. I also noted that printing just critic_param_gen won't do.
Any idea of why is that?
Also a similar thing happens if I use local_model = copy.deepcopy(global_model) as a function create_local_model, i.e. only works if that print is present.
In pseudo-code:
import torch.multiprocessiA3Cng as mp
import torch.nn as nn
import itertools as it
debug = True
A3C(nn.Module):
def __init__(self, model, n_features):
...
self.AC_architecture = SharedActorCritic(model, n_features)
class SharedActorCritic(nn.Module):
def __init__(self, model, n_features):
super(SharedActorCritic, self).__init__()
self.shared_architecture = model(n_features) # inherits from nn.Module
self.actor = SharedActor(n_features) # inherits from nn.Module
self.critic = SharedCritic(n_features) # inherits from nn.Module
self.critic_target = BaseCritic(model, n_features) # inherits from nn.Module
critic_param_gen = it.chain(self.shared_architecture.parameters(), self.critic.parameters())
print("checkpoint")
if debug: print(list(critic_param_gen)) # this makes the whole thing work
for trg_params, params in zip(self.critic_target.parameters(), critic_param_gen ):
trg_params.data.copy_(params.data)
def create_local_model(model, n_features):
local_model = A3C(model, n_features)
print("Process ended")
# in the main
global_model = Model() # works
global_model.share_memory() # doesn't really matter
p = mp.Process(target=create_local_model, args=(model, n_features, ))
p.start()
print("Process started")
p.join()
----
# output if debug is True
Process started
checkpoint
[ ...actual list of critic_param_gen ... ]
Process ended
# output if debug is False
Process started
checkpoint
# and then runs forever
Edit: solved the mystery about the print statement thanks to snakecharmerb. I created a minimal reproducible example. It seems that if the network is large enough, the copy operation breaks if executed in a process, but not outside of it (since global model can be instantiated).
import torch.nn as nn
import torch.multiprocessing as mp
import copy
class Net(nn.Module):
def __init__(self, n_features=256, n_layers=8):
super(Net, self).__init__()
self.net1 = nn.Sequential(*nn.ModuleList([nn.Linear(n_features, n_features) for _ in range(n_layers)]))
self.net2 = nn.Sequential(*nn.ModuleList([nn.Linear(n_features, n_features) for _ in range(n_layers)]))
for p1, p2 in zip(self.net1.parameters(), self.net2.parameters()):
p1.data.copy_(p2.data)
def forward(self, x):
return self.net(x)
def create_local_model_v1(global_model):
local_model = copy.deepcopy(global_model)
print("Process ended")
%%time
global_model = Net(16,2)
print("Global model created")
p = mp.Process(target=create_local_model_v1, args=(global_model,))
p.start()
print("Process started")
p.join()
# Output
Global model created
Process ended
Process started
CPU times: user 3 ms, sys: 11.9 ms, total: 14.9 ms
Wall time: 45.1 ms
%%time
global_model = Net(256,8)
print("Global model created")
p = mp.Process(target=create_local_model_v1, args=(global_model,))
p.start()
print("Process started")
p.join()
# Output - Gets stuck
Global model created
Process started