PyTorch Lightning DDP单节点多GPU训练：如何获取进程GPU分配信息

方法1：用PyTorch内置API直接获取 #

在DataLoader对应的Dataset的__getitem__方法里直接调用：

import torch

current_gpu_idx = torch.cuda.current_device()

DDP模式下每个进程会绑定到单独的GPU，这个调用会直接返回当前进程使用的GPU索引（0或1，对应你的两块GPU），拿来做分支处理就行。

方法2：读取自动设置的环境变量 #

单节点DDP中，PyTorch Lightning会自动设置LOCAL_RANK环境变量，其值就是当前进程对应的GPU索引，在Dataset里直接读：

import os

current_gpu_idx = int(os.environ.get("LOCAL_RANK", 0))

不用手动配置，直接用就行。

方法3：从LightningModule把GPU信息传给Dataset #

如果需要在初始化Dataset时就用到GPU信息，在LightningModule的dataloader方法里先拿到rank，再传给Dataset：

import pytorch_lightning as pl
from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __init__(self, gpu_idx):
        self.gpu_idx = gpu_idx
        # 这里根据gpu_idx做不同的初始化逻辑

class MyModel(pl.LightningModule):
    def train_dataloader(self):
        # 获取当前进程的local rank，对应GPU索引
        current_gpu_idx = self.trainer.local_rank
        dataset = MyDataset(gpu_idx=current_gpu_idx)
        return DataLoader(dataset)

单节点场景下global_rank和local_rank值一致，用self.trainer.global_rank也能得到相同结果。

注意事项 #

• 每个DDP进程会独立启动DataLoader的worker（如果设置了num_workers），worker会继承父进程的GPU绑定信息或环境变量，所以上述方法在worker里也能正常生效。

内容的提问来源于stack exchange，提问作者Moshel