【技术总结】Pytorch 多卡训练例程解析_多机多卡训练pytorch-CSDN博客

本文链接：https://blog.csdn.net/S_Kross/article/details/128525499

该代码展示了如何在PyTorch中使用DDP（DistributedDataParallel）进行多GPU训练，确保所有GPU的显存消耗一致。它涉及设置环境变量、初始化进程组、调整模型到指定设备以及使用DistributedSampler。训练过程中，模型在每个epoch后更新sampler，并且只在rank=0的进程中保存日志和模型。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

例程

import os
import pathlib
import numpy as np
import tqdm
import argparse
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import AdamW
from torch.optim.lr_scheduler import StepLR
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from utils import load_configs
from data import data_pipeline, get_ood_dataloader, RESIZE, CROP
from model import get_model
from ood_matrics import ACC, FPR, AUROC
from test import test



def setup(rank, world_size):
    # rank: 进程 ID 0-(world_size - 1)
    # world_size: 进程个数
    print("multi gpu setup", rank, world_size)
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12345'

    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    return 

def cleanup():
    dist.destroy_process_group()


def train(model: torch.nn.Module, ID_dataloader_train, ID_dataloader_val, OOD_dataloader, epoch, lr, checkpoint_path, fpr_threshold=0.95, log_path='./logs/baseline', device='cuda', current_rank=None, trainlog_writer=None):
    
    
    params = [
            {
   
   'params': model.od_head.parameters(), 'lr': lr},
            {
   
   'params': model.backbone.fc.parameters()},
            # {'params': model.backbone.layer4.parameters(), 'lr': 0.1 * lr}
    ]

    if