Hi, I’m currently using torch.multiprocessing for sending the outputs of a neural network to another process. For simple discussion, I have two processes: the first one is for loading training data, forwarding network and sending the results to the other one, while the other one is for recving the results from the previous process and handling the results.
The question I get troubled in is that, multiprocesses slow down the data loader. At every begining of each epoch, the forwarding process will unexpectedly get stucked for a while and then continue forwarding. When I cancel the multiprocess (i.e. just forward network and don’t send data across processes), all things go smoothly, and the stucked phenomenon at every begining epoch disappears.
Below is a python3 simple script (runs on one GTX 1080) which loads a pretrained network and use the dataset of Cifar10 with arguments of  num_workers=2 and pin_memory=True. You will observe the program gets stucked between for epoch in range(total_epoch) and for step, datas in enumerate(trainloader)
Anyone has any idea on why this happens and how to solve this? (I am doubting it is caused by limits of shared memory, but still don’t know how to avoid this)
import torch.multiprocessing as mp

def main():

  os.environ['CUDA_VISIBLE_DEVICES'] = '0'
  device = torch.device('cuda:0')

  model_name = 'vgg16'
  dataset = 'cifar10'

  total_epoch = 20
  batch_size = 128

  # get model and datasets
  model = get_a_pretrained_model
  trainloader = get_cifar10_trainloader # numworkers=2, pin_memory=True
  for param in model.parameters():
    param.requires_grad = False
  model = model.to(device)
  model.eval()

  mp.set_start_method('spawn')
  buffer = mp.Queue()
  recv_proc = mp.Process(target=recv, args=(buffer,))
  recv_proc.start()

  # forwarding model and send results
  for epoch in range(total_epoch):
    # program will get stucked here
    for step, datas in enumerate(trainloader):
      inputs, labels = datas
      inputs, labels = inputs.to(device), labels.to(device)
      outs = model(inputs)
      buffer.put(outs)
      if step % 20 == 0:
        print('Epoch [{}]\t Step [{}]'.format(epoch, step))
  buffer.put(None)

  recv_proc.join()
  

def recv(buffer):
  while True:
    recved = buffer.get()
    if recved is None:
      break
    # do something with recved
    del recved


if __name__ == '__main__':
  main()