DDP多GPU训练模型

训练代码采用了DDP，并是用torchrun来保证训练过程异常退出时，能够根据保存的模型接着训练。
训练代码：

import cifar10DataLoader as datasets
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader
import torchvision
import torch
import copy
import time
from tqdm import tqdm
from resnet import ResNetBase
from config import Config
from torch import nn, optim 

import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import os

# 初始化DDP
def ddp_setup():
    init_process_group(backend="nccl")

# 多GPU训练
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import os
# device = torch.device("cuda")
class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_data: DataLoader,
        test_data: DataLoader,
        optimizer: torch.optim.Optimizer,
        save_every: int,
        snapshot_path: str
    ) -> None:
        # 单机多GPU时时是用
        self.lock_rank = int(os.environ["LOCAL_RANK"])
        # 多机多GPU是用
        # self.global_rank = int(os.environ["RANK"])
        # 把模型放到GPU上
        self.module = model.to(self.lock_rank)
        self.train_data = train_data
        self.test_data = test_data
        self.optimizer = optimizer
        self.save_every = save_every
        self.epochs_run = 0
        if os.path.exists(snapshot_path):
            print("Loading snapshot")
            self._load_snapshot(snapshot_path)
        # 将模型交给DDP进行管理
        self.module = DDP(self.module, device_ids=[self.lock_rank])
    
    # 训练异常退出时，调用此函数接着训练之前的模型
    def _load_snapshot(self, snapshot_path):
        snapshot = torch.load(snapshot_path)
        self.module.load_state_dict(snapshot["MODEL_STATE"])
        self.epochs_run = snapshot["EPOCHS_RUN"]
        print(f"Resuming training from snapshot at Epoch {self.epochs_run}")

    
    def _run_batch(self, source, targets):
        self.optimizer.zero_grad()
        output = self.module(source)
        loss = torch.nn.CrossEntropyLoss()(output, targets)
        loss.backward()
        # loss.to(device)
        self.optimizer.step()
    
    def _run_epoch(self, epoch):
        b_sz = len(next(iter(self.train_data))[0])
        print(f'[GPU{self.lock_rank}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}')
        for source, targets in tqdm(self.train_data):
            # 将训练图片与标签均放入GPU中
            source = source.to(self.lock_rank)
            targets = targets.to(self.lock_rank)
            self._run_batch(source, targets)

    def _save_snapshot(self, epoch):
        # 使用了DDP
        snapshot = {}
        # DDP的模型保存在self.module.module.state_dict()中
        snapshot["MODEL_STATE"] = copy.deepcopy(self.module.module.state_dict())
        snapshot["EPOCHS_RUN"] = epoch
        torch.save(snapshot, 'snapshot.pt')
        print(f'Epoch {epoch} | Training checkpoint saved at snapshot.pt')

    def train(self, max_epochs: int):
        for epoch in range(self.epochs_run, max_epochs):
            self._run_epoch(epoch)
            if self.lock_rank == 0 and epoch % self.save_every == 0:
                self._save_snapshot(epoch)

def prepare_trainData(batch_size: int):
    dataset = datasets.CIFAR10_IMG('./data', train=True, transform=transforms.Compose([
        transforms.ToTensor(),
        # 先四周填充0，在把图像随机裁剪成32*32
        transforms.RandomCrop(32, padding=4),
        # 以0.5的概率水平翻转图片
        transforms.RandomHorizontalFlip(p=0.5),
        # 均值，标准差
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]))
    # 使用了DDP，需要修改sampler,同时将shuffle设为False
    return DataLoader(
        dataset, 
        batch_size=batch_size, 
        pin_memory=True, 
        shuffle=False,
        sampler=DistributedSampler(dataset))

def prepare_testData(batch_size: int):
    dataset = datasets.CIFAR10_IMG('./data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]))
    return DataLoader(
        dataset, 
        batch_size=batch_size, 
        pin_memory=True, 
        shuffle=False,
        sampler=DistributedSampler(dataset))
    
    
def load_train_objs():
    base = ResNetBase(Config.n_blocks, Config.n_channels, Config.bottlenecks, Config.first_kernel_size)
    classification = nn.Linear( Config.n_channels[-1], 10 )
    model = nn.Sequential( base, classification )
    # model = model.to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
    return model, optimizer

def main(save_every: int, total_epochs: int, snapshot_path: str="snapshot.pt"):
    # 调用函数对DDP进行初始化
    ddp_setup()
    model, optimizer = load_train_objs()
    train_data = prepare_trainData(batch_size=256)
    test_data = prepare_testData(batch_size=256)
    trainer = Trainer(model, train_data, test_data, optimizer, save_every, snapshot_path)
    trainer.train(total_epochs)
    destroy_process_group()


if __name__ == '__main__':
    import sys
    total_epochs = int(sys.argv[1])
    save_every = int(sys.argv[2])
    world_size = torch.cuda.device_count()
    print(world_size)
    # torch.cuda.set_device(world_size)
    main(save_every, total_epochs)    

#  torchrun --standalone --nproc_per_node=gpu cifarDDPnew.py 50 5 

# 此命令可以指定哪些GPU进行训练
# CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node=gpu cifarDDPnew.py 100 5

模型代码：

from typing import List, Optional

import torch
from torch import nn 
from typing import Any, TypeVar, Iterator, Iterable, Generic
class Module(torch.nn.Module):
    r"""
    Wraps ``torch.nn.Module`` to overload ``__call__`` instead of
    ``forward`` for better type checking.
    
    `PyTorch Github issue for clarification <https://github.com/pytorch/pytorch/issues/44605>`_
    """

    def _forward_unimplemented(self, *input: Any) -> None:
        # To stop PyTorch from giving abstract methods warning
        pass

    def __init_subclass__(cls, **kwargs):
        if cls.__dict__.get('__call__', None) is None:
            return

        setattr(cls, 'forward', cls.__dict__['__call__'])
        delattr(cls, '__call__')

    @property
    def device(self):
        params = self.parameters()
        try:
            sample_param = next(params)
            return sample_param.device
        except StopIteration:
            raise RuntimeError(f"Unable to determine"
                               f" device of {self.__class__.__name__}") from None


class ShortcutProjection(Module):
    def __init__(self, in_channels: int, out_channels:int, stride:int):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride)
        self.bn = nn.BatchNorm2d(out_channels)
    
    def forward(self, x:torch.Tensor):
        return self.bn(self.conv(x))
    
class ResidualBlock(Module):
    def __init__(self, in_channels: int, out_channels: int, stride: int):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)

        self.act1 = nn.ReLU()

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

        if stride != 1 or in_channels != out_channels:
            self.shortcut = ShortcutProjection(in_channels, out_channels, stride)
        else:
            self.shortcut = nn.Identity()
        
        self.act2 = nn.ReLU()

    def forward(self, x: torch.Tensor):
        shortcut = self.shortcut(x)
        x = self.act1(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        return self.act2(x + shortcut)

class BottleneckResidualBlock(Module):
    def __init__(self, in_channels: int, bottleneck_channels: int, out_channels: int, stride: int):
        super().__init__()

        self.conv1 = nn.Conv2d(in_channels, bottleneck_channels, kernel_size=1, stride=1)
        self.bn1 = nn.BatchNorm2d(bottleneck_channels)
        self.act1 = nn.ReLU()
        self.conv2 = nn.Conv2d(bottleneck_channels, bottleneck_channels, kernel_size=3, stride=stride, padding=1)
        self.bn2 = nn.BatchNorm2d(bottleneck_channels)
        self.act2 = nn.ReLU()
        self.conv3 = nn.Conv2d(bottleneck_channels, out_channels, kernel_size=1, stride=1)
        self.bn3 = nn.BatchNorm2d(out_channels)
        if stride != 1 or in_channels != out_channels:
            self.shortcut = ShortcutProjection(in_channels, out_channels, stride)
        else:
            self.shortcut = nn.Identity()
        self.act3 = nn.ReLU()
    
    def forward(self, x:torch.Tensor):
        shortcut = self.shortcut(x)
        x = self.act1(self.bn1(self.conv1(x)))
        x = self.act2(self.bn2(self.conv2(x)))
        x = self.bn3(self.conv3(x))
        return self.act3(x + shortcut)
    

class ResNetBase(Module):
    def __init__(self, n_blocks:List[int], n_channels: List[int],
                 bottelnecks: Optional[List[int]] = None,
                 img_channels: int = 3, first_kernel_size: int = 7):
        super().__init__()
        assert len(n_blocks) == len(n_channels)
        assert bottelnecks is None or len(bottelnecks) == len(n_channels)

        # // 向小取整
        self.conv = nn.Conv2d(img_channels, n_channels[0], 
                              kernel_size=first_kernel_size, stride=2, padding=first_kernel_size // 2)
        
        self.bn = nn.BatchNorm2d(n_channels[0])

        blocks = []

        prev_channels = n_channels[0]

        # enumerate枚举索引与值
        for i, channels in enumerate(n_channels):
            # 第一个stride为2，其他的为1
            stride = 2 if len(blocks) == 0 else 1

            if bottelnecks is None:
                # 不需要使用bottelnecks时，ResidualBlock：[3, 3]
                blocks.append(ResidualBlock(prev_channels, channels, stride=stride))
            else:
                # 后面的为BottleneckResidualBlock：[1, 3, 1]
                blocks.append(BottleneckResidualBlock(prev_channels, bottelnecks[i], channels,
                                                      stride=stride))
            prev_channels = channels

            # 需要多少个blocks
            for _ in range(n_blocks[i]-1):
                if bottelnecks is None:
                    blocks.append(ResidualBlock(channels, channels, stride=1))
                else:
                    blocks.append(BottleneckResidualBlock(channels, bottelnecks[i], channels, stride=1))
        
        self.blocks = nn.Sequential(*blocks)
    
    def forward(self, x:torch.Tensor):
        x = self.bn(self.conv(x))
        x = self.blocks(x)
        x = x.view(x.shape[0], x.shape[1], -1)
        return x.mean(dim=-1)

网络配置代码：

 class Config:
    n_channels = [16, 32, 64]
    bottlenecks = [8, 16, 16]
    n_blocks = [6, 6, 6]
    first_kernel_size = 3
    total_epoches = 500
    batch_size = 256
    Lr = 0.1

数据集代码：

import json
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import Dataset, DataLoader
from typing import Any, Callable, Optional, Tuple
# 继承Dataset类
class CIFAR10_IMG(Dataset):

    def __init__(self, root, train=True, transform = None, target_transform = None) -> None:
        super().__init__()
        self.train = train
        self.transform = transform
        self.target_transform = target_transform
        # 加载训练集
        if self.train:
            file_annotation = root + '/annotations/cifar10_train.json'
            img_folder = root + '/train_cifar10/'
        # 加载测试集
        else:
            file_annotation = root + '/annotations/cifar10_test.json'
            img_folder = root + '/test_cifar10/'
        # 读取json文件
        fp = open(file_annotation, 'r')
        data_dict = json.load(fp)
        # 图片数和标签数不匹配说明数据集标注有问题，报错
        assert len(data_dict['images'])==len(data_dict['categories'])
        num_data = len(data_dict['images'])
        # 读取图片与对应的标注
        self.filenames = []
        self.labels = []
        self.img_folder = img_folder
        for i in range(num_data):
            self.filenames.append(data_dict['images'][i])
            self.labels.append(data_dict['categories'][i])

    def __getitem__(self, index):
        img_name = self.img_folder + self.filenames[index]
        label = self.labels[index]
        # 将数据转换为numpy格式
        img = plt.imread(img_name)
        if self.transform is not None:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.filenames)