Как использовать многопроцессорность в PyTorch?

Question

Я пытаюсь использовать PyTorch со сложной функцией потерь.Чтобы ускорить код, я надеюсь, что смогу использовать пакет многопроцессорной обработки PyTorch.

В первой пробной версии я поместил функции 10x1 в NN и получил вывод 10x4.

После этого яхочу передать 10x4 параметров в функцию, чтобы сделать некоторые вычисления.(Расчет будет сложным в будущем.)

После расчета функция вернет массив 10x1.Этот массив будет установлен как NN_energy и будет вычислять функцию потерь.

Кроме того, я также хочу знать, есть ли другой метод для создания обратно-способного массива для хранения массива NN_energy вместо использования

NN_energy = net(Data_in)[0:10,0]

Большое спасибо.

Полный код:

import torch
import numpy as np
from torch.autograd import Variable 
from torch import multiprocessing

def func(msg,BOP):
    ans = (BOP[msg][0]+BOP[msg][1]/BOP[msg][2])*BOP[msg][3]
    return ans

class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden_1, n_hidden_2, n_output):
        super(Net, self).__init__()
        self.hidden_1 = torch.nn.Linear(n_feature , n_hidden_1)  # hidden layer
        self.hidden_2 = torch.nn.Linear(n_hidden_1, n_hidden_2)  # hidden layer
        self.predict  = torch.nn.Linear(n_hidden_2, n_output  )  # output layer

    def forward(self, x):
        x = torch.tanh(self.hidden_1(x))      # activation function for hidden layer
        x = torch.tanh(self.hidden_2(x))      # activation function for hidden layer
        x = self.predict(x)                   # linear output
        return x

if __name__ == '__main__': # apply_async
    Data_in      = Variable( torch.from_numpy( np.asarray(list(range( 0,10))).reshape(10,1) ).float() )
    Ground_truth = Variable( torch.from_numpy( np.asarray(list(range(20,30))).reshape(10,1) ).float() )

    net = Net( n_feature=1 , n_hidden_1=15 , n_hidden_2=15 , n_output=4 )     # define the network
    optimizer = torch.optim.Rprop( net.parameters() )
    loss_func = torch.nn.MSELoss()  # this is for regression mean squared loss 
    NN_output = net(Data_in)   
    args = range(0,10)
    pool = multiprocessing.Pool()
    return_data = pool.map( func, zip(args, NN_output) )
    pool.close()
    pool.join()

    NN_energy = net(Data_in)[0:10,0]  
    for i in range(0,10):
        NN_energy[i] = return_data[i]

    loss = torch.sqrt( loss_func( NN_energy , Ground_truth ) )     # must be (1. nn output, 2. target) 
    print(loss)

Сообщения об ошибках:

Файл "C: \ ProgramData \ Anaconda3 \ lib \ site-packages \ torch \ multiprocessing \ reductions.py", строка 126, в lower_tensor повысить RuntimeError ("трусливый отказ сериализации неконечного тензора, который требует_grad,"

RuntimeError: трусливый отказ от сериализации неконечного тензора, который требует _grad, так как autograd не поддерживает пересечение границ процесса. Если вы просто хотите передать данные, вызовите detach () для тензора перед сериализацией (например, поместите его в очередь).