import os import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader, DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP import torch.distributed as dist from torch.multiprocessing import Process import matplotlib.pyplot as plt class DenseBlock(nn.Module): def __init__(self, in_channels, growth_rate, num_layers): super(DenseBlock, self).__init__() self.layers = nn.ModuleList() for i in range(num_layers): layer = self._make_layer(in_channels + i * growth_rate, growth_rate) self.layers.append(layer) def _make_layer(self, in_channels, growth_rate): return nn.Sequential( nn.BatchNorm2d(in_channels), nn.ReLU(inplace=True), nn.Conv2d(in_channels, 4 * growth_rate, kernel_size=1, bias=False), nn.BatchNorm2d(4 * growth_rate), nn.ReLU(inplace=True), nn.Conv2d(4 * growth_rate, growth_rate, kernel_size=3, padding=1, bias=False) ) def forward(self, x): features = [x] for layer in self.layers: new_feature = layer(torch.cat(features, 1)) features.append(new_feature) return torch.cat(features, 1) class TransitionLayer(nn.Module): def __init__(self, in_channels, out_channels): super(TransitionLayer, self).__init__() self.transition = nn.Sequential( nn.BatchNorm2d(in_channels), nn.ReLU(inplace=True), nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), nn.AvgPool2d(kernel_size=2, stride=2) ) def forward(self, x): return self.transition(x) class DenseNet(nn.Module): def __init__(self, growth_rate=12, block_config=(6, 12, 24, 16), num_init_features=64, num_classes=10): super(DenseNet, self).__init__() # 初始卷积层 self.features = nn.Sequential( nn.Conv2d(3, num_init_features, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(num_init_features), nn.ReLU(inplace=True) ) # DenseBlocks 和 TransitionLayers num_features = num_init_features self.dense_blocks = nn.ModuleList() self.transitions = nn.ModuleList() for i, num_layers in enumerate(block_config): # 添加 DenseBlock block = DenseBlock(num_features, growth_rate, num_layers) self.dense_blocks.append(block) num_features += num_layers * growth_rate # 添加 TransitionLayer (除了最后一个block) if i != len(block_config) - 1: transition = TransitionLayer(num_features, num_features // 2) self.transitions.append(transition) num_features = num_features // 2 # 最终分类层 self.classifier = nn.Sequential( nn.BatchNorm2d(num_features), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(num_features, num_classes) ) def forward(self, x): x = self.features(x) for i, block in enumerate(self.dense_blocks): x = block(x) if i < len(self.transitions): x = self.transitions[i](x) x = self.classifier(x) return x def get_cifar10_dataloader(batch_size=128, distributed=False, rank=0, world_size=1): # 数据预处理 transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) # 加载数据集 trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) if distributed: train_sampler = DistributedSampler(trainset, num_replicas=world_size, rank=rank) trainloader = DataLoader(trainset, batch_size=batch_size, sampler=train_sampler, num_workers=2) else: trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2) testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2) return trainloader, testloader def train_epoch(model, trainloader, criterion, optimizer, device): model.train() running_loss = 0.0 correct = 0 total = 0 for batch_idx, (inputs, targets) in enumerate(trainloader): inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() running_loss += loss.item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() return 100. * correct / total def test(model, testloader, criterion, device): model.eval() test_loss = 0 correct = 0 total = 0 with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(testloader): inputs, targets = inputs.to(device), targets.to(device) outputs = model(inputs) loss = criterion(outputs, targets) test_loss += loss.item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() acc = 100. * correct / total return acc def main(rank, world_size): # 初始化分布式环境 dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=world_size) torch.cuda.set_device(rank) device = torch.device(f'cuda:{rank}') # 加载数据 trainloader, testloader = get_cifar10_dataloader(batch_size=128, distributed=True, rank=rank, world_size=world_size) # 创建模型 model = DenseNet(growth_rate=12, block_config=(6, 12, 24, 16), num_init_features=24, num_classes=10).to(device) model = DDP(model, device_ids=[rank]) # 损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4) scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1) # 训练 num_epochs = 20 for epoch in range(num_epochs): print(f'Rank {rank}, Epoch [{epoch+1}/{num_epochs}]') # 训练一个epoch train_acc = train_epoch(model, trainloader, criterion, optimizer, device) # 测试 if rank == 0: # 只在主进程打印测试结果 test_acc = test(model, testloader, criterion, device) print(f'Test Accuracy: {test_acc:.2f}%') # 更新学习率 scheduler.step() dist.destroy_process_group() if __name__ == '__main__': world_size = torch.cuda.device_count() os.environ['MASTER_ADDR'] = '127.0.0.1' os.environ['MASTER_PORT'] = '29500' processes = [] for rank in range(world_size): p = Process(target=main, args=(rank, world_size)) p.start() processes.append(p) for p in processes: p.join()