multiprocessing pytorch 代码

下面是一个使用multiprocessing库在PyTorch中进行多进程训练的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

def train(model, optimizer, rank, input_data, target):
    optimizer.zero_grad()
    output = model(input_data)
    loss = nn.MSELoss()(output, target)
    loss.backward()
    optimizer.step()
    print('Rank', rank, 'loss:', loss.item())

if __name__ == '__main__':
    mp.set_start_method('spawn')
    num_processes = 4
    num_train_samples = 100
    input_size = 10
    output_size = 1
    learning_rate = 0.01
    num_epochs = 20
    batch_size = 10

    # Generate input and target data
    input_data = torch.randn(num_train_samples, input_size)
    target = torch.randn(num_train_samples, output_size)

    # Create model and optimizer
    model = Net()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)

    # Create a list of processes
    processes = []
    for rank in range(num_processes):
        p = mp.Process(target=train, args=(model, optimizer, rank, input_data, target))
        p.start()
        processes.append(p)

    # Wait for all processes to finish
    for p in processes:
        p.join()

    # Print final model parameters
    print(model.state_dict())

在这个示例中，我们定义了一个简单的神经网络模型（Net），并且使用多进程来训练模型。每个进程通过调用train函数来执行模型的前向传递、反向传递和优化步骤。我们使用PyTorch的优化器（SGD）和损失函数（MSELoss）来训练模型。在训练过程中，我们打印出每个进程的损失值，并且在所有进程完成训练后打印出最终的模型参数。

要注意的是，为了在Windows平台上使用多进程，我们需要在脚本的开头调用mp.set_start_method('spawn')来设置启动方法。在Linux和MacOS上，默认启动方法是'fork'，但是在Windows上不支持该启动方法。