如何在PyTorch中用两块不同GPU同时训练两个模型（同数据）

Hey there! Let's get your two models running simultaneously on separate GPUs using PyTorch's spawn method—this is exactly the kind of use case multiprocessing was built for. Here's a complete, actionable implementation tailored to your setup:

Step-by-Step Implementation #

First, let's break down the core idea: we'll use torch.multiprocessing.spawn to launch two separate subprocesses, each bound to one of your GPUs. Each subprocess will handle moving its assigned model to the correct GPU, running the training/testing loop, and sending the accuracy result back to the main process.

Full Code Example #

import torch
import torch.multiprocessing as mp

# Assume your existing training function looks like this (adjust to match your actual code)
def training(model, train_dl, test_dl):
    # Your existing training logic here—ensure all tensor operations use the model's device
    model.train()
    for batch in train_dl:
        inputs, labels = batch
        inputs, labels = inputs.to(next(model.parameters()).device), labels.to(next(model.parameters()).device)
        # ... rest of your training steps (forward pass, loss calculation, optimizer step) ...
    
    # Test phase to compute accuracy
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch in test_dl:
            inputs, labels = batch
            inputs, labels = inputs.to(next(model.parameters()).device), labels.to(next(model.parameters()).device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

# Worker function: runs independently in each subprocess
def worker(rank, models, train_dl, test_dl, result_queue):
    # Bind this subprocess exclusively to the specified GPU
    torch.cuda.set_device(rank)
    
    # Grab the assigned model and move it to the current GPU
    model = models[rank]
    model = model.to(f"cuda:{rank}")
    
    # Execute training and get the accuracy
    accuracy = training(model, train_dl, test_dl)
    
    # Send the result back to the main process (rank lets us map to net1/net2)
    result_queue.put((rank, accuracy))

if __name__ == "__main__":
    # --- Replace these with your actual model and data loader definitions ---
    net1 = torch.nn.Linear(10, 2)  # Example model (swap with your net1)
    net2 = torch.nn.Linear(10, 2)  # Example model (swap with your net2)
    train_dl = torch.utils.data.DataLoader(...)  # Your existing training data loader
    test_dl = torch.utils.data.DataLoader(...)   # Your existing test data loader
    
    # Create a queue to safely pass results from subprocesses to main process
    result_queue = mp.Queue()
    
    # Launch 2 subprocesses (one for each GPU)
    mp.spawn(
        worker,
        args=( [net1, net2], train_dl, test_dl, result_queue ),
        nprocs=2,
        join=True  # Wait for all subprocesses to finish before proceeding
    )
    
    # Retrieve and organize the results
    results = {}
    while not result_queue.empty():
        rank, acc = result_queue.get()
        results[rank] = acc
    
    accuracy1 = results[0]
    accuracy2 = results[1]
    
    print(f"net1 test accuracy: {accuracy1:.4f}")
    print(f"net2 test accuracy: {accuracy2:.4f}")

Critical Tips to Avoid Headaches #

• Data Loader Gotchas: When using spawn, PyTorch switches to a different multiprocessing context. If your data loaders use num_workers > 0, you might hit serialization errors. Either set num_workers=0 or re-initialize your data loaders inside the worker function (instead of passing them from the main process) to fix this.
• Model Serialization: Ensure your custom models don't have non-serializable attributes (like open file handles or custom objects that can't be pickled)—this will break subprocess communication.
• GPU Isolation: torch.cuda.set_device(rank) guarantees each subprocess only uses its assigned GPU, preventing cross-GPU resource conflicts.
• Result Passing: Subprocesses created with spawn can't return values directly, so we use mp.Queue() to safely send accuracy results back to the main process.

How This Aligns With Your Requirements #

1. net1.to("cuda:0") → Handled automatically in the worker with rank 0
2. net2.to("cuda:1") → Handled automatically in the worker with rank 1
3. accuracy1 = training(net1, train_dl, test_dl) → Runs in subprocess 0, result sent back to main process
4. accuracy2 = training(net2, train_dl, test_dl) → Runs in subprocess 1, result sent back to main process

Both models train and test at the same time, fully utilizing your two GPUs instead of running sequentially!

内容的提问来源于stack exchange，提问作者Gustavo Vargas Hakim