深度学习疑问记录
1 损失值的计算
-
在分批训练中，计算每个 epoch 的损失通常涉及以下步骤：
- 初始化累积损失：在每个 epoch 的开头，设置一个变量 running_loss 为 0。
- 逐个批次处理：
- 前向传播：使用当前的模型参数计算预测输出。
- 计算损失：使用损失函数计算预测输出和真实标签之间的损失。
- 反向传播：执行反向传播，计算损失相对于模型参数的梯度。
- 参数更新：使用优化器更新模型参数。
- 累积损失：将当前批次的损失累加到 running_loss 中。
- 计算平均损失：在处理完所有批次后，计算平均损失，即 epoch_loss = running_loss / len(train_loader)。
-
注意点1： 平均损失的计算
PyTorch的损失函数（如 CrossEntropyLoss）默认对每个批次的损失取均值（mean），也即损失函数的参数reduction的默认值为reduction='mean'。
| 模式 | 行为 | 适用场景 |
|---|---|---|
'mean' |
计算批次内所有样本损失的平均值（默认值） | 多数分类/回归任务 |
'sum' |
直接返回批次内所有样本损失的总和 | 需要自定义加权时 |
'none' |
返回每个样本的独立损失值（未汇总的张量） | 需手动处理损失值时 |
- 注意点2： 累加的时候要注意使用
.item()
.item()的作用：该方法从单元素的张量中提取对应的 Python 标量值（如 float 或 int），并断开与计算图的连接
这个问题主要是关于损失值到底是求和还是求平均，暂时以下面的代码为标准：
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
# 超参数设置
BATCH_SIZE = 64
EPOCHS = 10
LR = 0.01
# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. 数据准备
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(
root='./data',
train=True,
download=True,
transform=transform
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
shuffle=True
)
# 2. 定义神经网络模型
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(28*28, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = x.view(-1, 28*28) # 展平图像
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
x = self.fc3(x)
return x
model = SimpleNet().to(device)
# 3. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=LR)
# 记录loss变化的列表
epoch_losses = []
# 4. 训练循环
for epoch in range(EPOCHS):
model.train()
running_loss = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
# 前向传播
outputs = model(data)
loss = criterion(outputs, target)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 累加loss
running_loss += loss.item()
# 计算epoch平均loss
epoch_loss = running_loss / len(train_loader)
epoch_losses.append(epoch_loss)
print(f'Epoch [{epoch+1}/{EPOCHS}], Loss: {epoch_loss:.4f}')
# 5. 绘制loss曲线
plt.figure(figsize=(10, 5))
plt.plot(range(1, EPOCHS+1), epoch_losses, marker='o')
plt.title('Training Loss Curve')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.grid(True)
plt.xticks(range(1, EPOCHS+1))
plt.show()
2 分布式训练
2.1 理论简单介绍
分布式训练总览：
下面主要介绍使用数据并行和DDP的方法。
关于图中用到的通信协议AllReduce(Ring AllReduce)的图示解说请参考链接1。
简单看一下DDP的后端通信方式：
- DDP的3种启动方法：
mp.spawn()启动：mp模块完成对multiprocessing库进行封装，并没有特定针对DDP。tochrnn启动：运行脚本的命令由python变为了torchrun。- 相较于使用
mp.spawn()启动，torchrun会自动控制一些环境变量的设置，只需设置os.envirn['CUDA_VISIBLE_DEVICES'](不设置默认为该机器上的所有GPU)，无需设置os:envirn['MASTER_ADDR']等环境变量，因而更为方便。
- 相较于使用
torch.distributed.launch启动:代码量更少，启动速度更快。(即将被淘汰)
DDP在使用前必须进行初始化，下面是两种初始化方法（推荐使用ENV方法，不使用Tcp）：
2.2 分步展示程序修改
下面分步骤展示如何将一个单机单卡的程序修改为单机多卡
- ① 相关库导入
# dist: 多卡通讯
import torch.distributed as dist
# mp: DDP程序启动
import torch.multiprocessing as mp
# Gradscaler: 混合精度训练
from torch.cuda.amp import GradScaler
# DistributedSampler: 数据采样
from torch.utils.data.distributed import DistributedSampler
# DDP: 模型传递
from torch.nn.parallel import DistributedDataParallel as DDP
- ② 程序修改--定义关键函数
# init_ddp(local_rank): 对进程进行初始化，使用 nccl 协议进行后端通信，并用 env 作为初始化方法
def init_ddp(local_rank):
# 这一句之后，在转换device的时候直接使用 a=a.cuda()即可，否则要用a=a.cuda(local_rank)
torch.cuda.set device(local_rank)
os.environ['RANK']= str(local_rank)
dist.init_process_group(backend='nccl', init_method='env://')
pass
local_rank = dist.get_rank() # 当前进程序号(GPU号)
world_size = dist.get_world_size() # 总进程数
# reduce_tensor(tensor): 对多个进程的计算结果进行汇总，如 loss、评价指标
def reduce_tensor(input_tensor):
"""
对多个进程计算的多个 tensor 类型的输出值取平均操作
"""
# 创建输入张量的副本，避免直接修改原始张量
rt = input_tensor.clone()
# 注意：clone()方法会复制张量的数据和计算图（保留梯度），但新张量与原始张量无内存共享
# 对所有进程中的 rt 张量执行全局求和，并将结果同步到所有进程
dist.all_reduce(rt, op=dist.reduce_op.SUM)
# 注意：dist.all_reduce()是PyTorch分布式通信（torch.distributed）的核心函数，要求所有进程同时调用此函数
# 执行后，所有进程的 rt 张量的值会被替换为全局求和的结果
# 将全局求和后的张量除以进程总数，得到全局平均值
rt /= dist.get world size()
return rt
# get_ddp_generator(seed): 用于训练过程中，增强训练的随机性
def get_ddp_generator(seed):
# 对每个进程使用不同的随机种子，增强训练的随机性
local_rank = dist.get_rank()
g = torch.Generator()
g.manual_seed(seed + local rank)
return g
- ③ 程序修改--程序入口
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-args', help="priority", type=bool, required=False, default=True)
# 下面这行gpu参数需要添加======================================================
# 当前机器有3个GPU，则定义0、1、2
parser.add_argument('-gpu', default='0,1,2', type=str, help='gpu device ids for CUDA_VISIBLE_DEVICES')
parser.add_argument('-mode', help="train&test", type=str, required=false, default='train')
parser.add_argument('-requires_grad', help="whether to weight decay", type=bool, required=false)
args = parser.parse_args()
# 0号机器的IP
os.environ['MASTER_ADDR']='localhost' # 0号机器的IP
# 0号机器的可用端口
os.environ['MASTER_PORT']='19198' # 0号机器的可用端口
# 使用哪些GPU
os.environ['CUDA_VISIBLE_DEVICES']= args['gpu']
# 总进程数 = 可用GPU数量(每个GPU看作一个进程)
world_size = torch.cuda.device_count()
# 环境变量赋值
os.environ['WORLD_SIZE']= str(world size)
if args['mode']== 'train':
time_start = time.time()
mp.spawn(fn=main, args=(args,), nprocs=world_size) # 注意：这里args必须写成元组的形势
time_elapsed =time.time()- time_start
print(f'\ntime elapsed:{time_elapsed:.2f} seconds.')
- ④ 程序修改--
main()函数
# main函数修改点1：参数列表更新，添加额外参数local_rank，该参数无需在mp.spawn()函数中传递，系统会自动分配
def main(local_rank, args):
# main函数修改点2： 进程初始化，调用 init_ddp() 函数实现
init_ddp(local_rank)
best_macro=0
model, tokenizer = initialise_model(args['modelname'], args['num_labels'])
# 模型放入设备
model.cuda()
# main函数修改点3：BN层同步，调用 convert_sync_batchnorm() 函数实现
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
num_gpus = torch.cuda.device_count()
if num_gpus>1:
print('use {} gpus!'.format(num_gpus))
# main函数修改点4：模型封装，调用 DistributedDataParallel() 函数实现
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
"""
中间代码省略
"""
# main函数修改点5：混合精度训练，调用 Gradscaler() 函数实现，作为参数传至 train() 函数中
scaler = GradScaler()
criterion = BCEWithLogitsLoss().cuda() # 定义损失函数并放到设备上
train_dataloader = get_dataloader(args['traincsvpath'], args, tokenizer, train-True
valid_dataloader = get_dataloader(args['valcsvpath'], args, tokenizer, train=False)
for actual_epoch in trange(args['num_epochs'], desc="Epoch"):
# main函数修改点6：避免副本重复执行，可以只允许rank0号进程打印
if local_rank == 0:
print("begin training of epoch %d / %d" % (actual_epoch + 1, args['num_epochs']))
# main函数修改点7：训练采样器设置，每个epoch设置不同的sampling顺序
train_dataloader.sampler.set_epoch(actual_epoch)
train(model, train_dataloader, optimizer, scheduler, criterion, actual epoch, scaler, args)
"""
中间代码省略
"""
# main函数修改点8：消除进程组，调用destroy_process_group()函数实现
dist.destroy_process_group()
- ⑤ 程序修改——
get_dataloader()函数
def get_dataloader(path, args, tokenizer, train):
"""
根据给定的路径获取数据，并将数据和训练标志传递给数据加器，这样可以方便地从给定路径加载数据并生成数据加载器，以供后续的模
path:数据存放路径
tokenizer:分词器
train:是否是训练阶段
"""
texts, labels = load dataset(path, args['num_labels'])
texts = tokenizer(texts, padding='max length', truncation=True, return tensors='pt', max_length=args[
data = TensorDataset(texts['input_ids'], texts['attention_mask'l, torch.tensor(labels))
if train:
# 创建一个随机采样器
train_sampler= DistributedSampler(data, shuffle=True)
# 使用 get_ddp_generator()函数增强随机性
g = get_ddp_generator()
dataloader = DataLoader(dataset=data,
batch size=args['batch_size'],
num_workers=args['num_workers'],
pin_memory=True,
shuffle=False, # 注意这里的shuffle关闭，因为下面使用了随机采样器
sampler=train_sampler, # 采用随机采样器
generator=g)
else:
# 创建一个顺序采样器
test_sampler = DistributedSampler(data, shuffle=False)
dataloader = DataLoader(dataset=data,
batch_size=args['batch_size'],
num_workers=args['num_workers'],
pin_memory=True,
shuffle=False,
sampler=test_sampler)
return dataloader
- ⑥ 程序修改——
train()、val函数
def train(model, train_dataloader, optimizer, scheduler, criterion, actual_epoch, scaler, args):
model.train()
tr_loss = 0
num_train_samples = 0
for step, batch in enumerate(train_dataloader):
batch = tuple(t.cuda(non_blocking=True) for t in batch)
b_input_ids, b_input_mask, b_labels = batch
# 前向传播阶段使用autocast控制半精度计算
with torch.cuda.amp.autocast():
output = model(b_input _ids, attention_mask=b_input_mask, labels=b_labels) # 运行到这一行会增加
loss = criterion(output.logits.view(-1, args['num_labels']), b_labels, type_as(output.logits).view
# 使用前面定义的函数reduce_tensor对并行进程计算多个loss取平均
reduced_loss = reduce_tensor(loss.data)
if dist.get_rank()==0:
print("\n output Loss: ", reduced_loss.item())
tr_loss += reduced_loss.item()
# 并行状态下的更新，不同进程分别根据自己计算的 loss 更新数据
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer) # 运行到这一行会增加一下显存
# 下面四行，多个进程只执行一次
scheduler.step()
scaler.update()
num_train_samples +=b_labels.size(0) # 将批次中的样本数量添加到 num_train_samples 中
torch.cuda.empty_cache() # 释放GPU reserved memory显存
epoch_train_loss = tr_loss / num_train_samples # num_train_samples 代表每个进程承接的样本数量，由于上面己经
if dist.get_rank()== 0:
print("\nTrain loss after Epach {} : {}".format(actual_epoch, epoch_train_loss))
@torch.no grad()
def validate(model, valid dataloader, criterion, epoch, args, threshold=0.5):
model.eval()
"""
省略部分代码
"""
with torch.no_grad():
with torch.cuda.amp.autocast()
output = model(b_input_ids, attention_mask=b_input_mask)
logits = output.logits
loss = criterion(logits.view(-1,args['num_labels']), b_labels.type_as(logits)
reduced_loss = reduce_tensor(loss.data)
eval_loss += reduced_loss.item()
if dist.get_rank() == 0:
print("Validation loss after Epoch {} : {}".format(epoch, epoch_eval_loss))
# 每个并行进程都会分别执行下列计算操作，得到各进程对应的macro评价指标
pred_labels = [item for sublist in pred_labels for item in sublist]
true_labels = [item for sublist in true labels for item in sublist]
pred_bools = [pl>threshold for pl in pred labels]
true_bools = [tl==1 for tl in true labels]
macro = f1_score(true bools, pred_bools, average='macro')
# 汇总不同实验结果
macro = redce_tensor(torch.tensor(macro).cuda())
return macro
下面给出一个完整的分类网络的分布式训练代码，代码来自GitHub仓库：rickyang1114-DDP-practice
import os
import time
import argparse
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
# 辅助完成不同进程(GPU)通信
import torch.distributed as dist
"""
辅助完成启动多进程，是多进程管理工具，为DDP等分布式方案提供底层进程调度支持
底层进程管理：提供spawn等接口启动子进程，为DDP的多进程模式提供底层支持（如spawn接口可自动启动N个进程，每个进程对应一个GPU）。
不直接参与训练逻辑：仅负责进程创建和通信上下文初始化，具体的分布式训练逻辑（如梯度同步）由DDP或Horovod实现。
"""
import torch.multiprocessing as mp
# 辅助完成混合精度训练
from torch.cuda.amp import GradScaler
class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.fc = nn.Linear(7 * 7 * 32, num_classes)
def forward(self, x):
# utilize mixed precision training to accelerate training and inference
with torch.cuda.amp.autocast(): # 启用混合精度训练
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.fc(out)
return out
def prepare():
parser = argparse.ArgumentParser()
parser.add_argument("--gpu", default="0,1") # 指定使用的GPU ID
parser.add_argument("-e", "--epochs", default=3, type=int, metavar="N", help="number of total epochs to run") # 指定训练轮数
parser.add_argument("-b", "--batch_size", default=32, type=int, metavar="N", help="number of batchsize") # 批量大小
args = parser.parse_args()
# 设置分布式训练所需的环境变量
# The following environment variables are set to enable DDP
os.environ["MASTER_ADDR"] = "localhost" # 主进程(0号GPU)的IP
os.environ["MASTER_PORT"] = "19198" # 主进程(0号GPU)的可用端口号
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # 指定具体使用哪些GPU
world_size = torch.cuda.device_count()
os.environ["WORLD_SIZE"] = str(world_size)
return args
def init_ddp(local_rank):
"""
进程组的初始化
参数：
local_rank: 当前进程在本地机器的GPU索引（如0或1或2...）
"""
# after this setup, tensors can be moved to GPU via `a = a.cuda()` rather than `a = a.to(local_rank)`
torch.cuda.set_device(local_rank) # 设置当前进程使用的GPU
os.environ["RANK"] = str(local_rank)
# 初始化分布式后端（nccl用于GPU通信，初始化模式为env），通过环境变量自动读取配置
dist.init_process_group(backend="nccl", init_method="env://")
def get_ddp_generator(seed=3407):
"""
辅助增加训练的随机性，主要在train_dataloader中调用
"""
local_rank = dist.get_rank() # 获取当前进程的全局ID
g = torch.Generator()
g.manual_seed(seed + local_rank) # 为每个进程设置不同种子，确保数据分布一致
return g
def train(model, train_dloader, criterion, optimizer, scaler):
# 设置模型为训练模式
model.train()
for images, labels in train_dloader:
# 将数据移动到GPU
images = images.cuda()
labels = labels.cuda()
outputs = model(images)
loss = criterion(outputs, labels)
# 清空梯度
optimizer.zero_grad()
# 混合精度反向传播
scaler.scale(loss).backward()
"""
说明：在分布式训练中，每个GPU（进程）处理数据的一部分，计算自己的本地损失。
PyTorch的DDP模块会在调用loss.backward()时自动执行梯度同步，通过all-reduce操作
将所有进程的梯度求和并平均，确保所有进程的梯度一致。故无需手动对损失进行求和或平均。
但是：如果需要在日志中输出全局平均loss（如统计训练损失），需手动对各个GPU的loss取平均
reduced_loss = reduce_tensor(loss) # 使用自定义reduce_tensor函数
"""
# 使用优化器更新模型参数
scaler.step(optimizer)
# 更新缩放器
scaler.update() # 调整梯度缩放因子，为下一次迭代做准备，缩放因子会根据当前梯度的动态范围自动调整
# 使用前面定义的函数reduce_tensor对并行进程计算多个loss取平均
reduced_loss = reduce_tensor(loss.data)
if dist.get_rank()==0:
print("\n output Loss: ", reduced_loss.item())
tr_loss += reduced_loss.item()
def test(model, test_dloader):
local_rank = dist.get_rank()
# 设置模型为评估模式
model.eval()
size = torch.tensor(0.0).cuda()
correct = torch.tensor(0.0).cuda()
for images, labels in test_dloader:
images = images.cuda()
labels = labels.cuda()
with torch.no_grad():
outputs = model(images)
size += images.size(0) # 累加当前batch的样本数
# 累加当前batch中预测正确的样本数
correct += (outputs.argmax(1) == labels).type(torch.float).sum()
# 将各进程的size汇总到主进程(0号GPU、rank 0)
dist.reduce(size, 0, op=dist.ReduceOp.SUM)
# # 将各进程的correct汇总到rank 0
dist.reduce(correct, 0, op=dist.ReduceOp.SUM)
if local_rank == 0:
acc = correct / size
print(f"Accuracy is {acc:.2%}")
def main(local_rank, args):
# 初始化DDP
init_ddp(local_rank) # 若有3个GPU，则此初始化函数运行3次，local_rank=0、local_rank=1、local_rank=2
# 模型移至GPU
model = ConvNet().cuda() # Note: the `forward` method of the model has been modified
# 转换为同步批归一化
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
# 包装模型以支持分布式训练
model = nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank]
) ### Wrap with DDP
# 定义损失函数并放入设备
criterion = nn.CrossEntropyLoss().cuda()
# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), 1e-4)
# 混合精度缩放器
scaler = GradScaler() # Used for mixed precision training
# 数据加载
train_dataset = torchvision.datasets.MNIST(
root="./data", train=True, transform=transforms.ToTensor(), download=True
)
# 用于DDP(分布式)模式的采样器
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) # Sampler specifically for DDP
g = get_ddp_generator()
train_dloader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=args.batch_size,
shuffle=False, ### shuffle is mutually exclusive with sampler
num_workers=4,
pin_memory=True,
sampler=train_sampler,
generator=g,
) # generator is used for random seed
test_dataset = torchvision.datasets.MNIST(
root="./data", train=False, transform=transforms.ToTensor(), download=True
)
test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset) # Sampler specifically for DDP
test_dloader = torch.utils.data.DataLoader(
dataset=test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=test_sampler,
)
for epoch in range(args.epochs):
if local_rank == 0: # avoid redundant printing for each process
print(f"begin training of epoch {epoch + 1}/{args.epochs}")
# 为每个epoch重新打乱数据，即每个epoch的数据采样顺序都不同
train_dloader.sampler.set_epoch(epoch) # set epoch for sampler
train(model, train_dloader, criterion, optimizer, scaler)
if local_rank == 0:
print(f"begin testing")
test(model, test_dloader)
if local_rank == 0: ### avoid redundant saving for each process
torch.save(
{"model": model.state_dict(), "scaler": scaler.state_dict()},
"ddp_checkpoint.pt",
)
# 销毁进程组
dist.destroy_process_group() ### destroy the process group, in accordance with init_process_group.
if __name__ == "__main__":
args = prepare()
time_start = time.time()
mp.spawn(main, args=(args,), nprocs=torch.cuda.device_count())
time_elapsed = time.time() - time_start
print(f"\ntime elapsed: {time_elapsed:.2f} seconds")
最后，再给出一个在kaggle上使用两个T4 GPU训练的代码：
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, DistributedSampler
import os
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group("nccl", rank=rank, world_size=world_size)
model = nn.Sequential(
nn.Linear(784, 4096),
nn.ReLU(),
nn.Linear(4096, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10)
).to(rank)
ddp_model = DDP(model, device_ids=[rank])
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
loader = DataLoader(dataset, batch_size=128, sampler=sampler)
for epoch in range(200):
print('当前正在进行第{}轮训练'.format(epoch))
if epoch == 1:
print('当前模型所在设备: ', next(model.parameters()).device)
pass
sampler.set_epoch(epoch)
for data, target in loader:
data, target = data.to(rank), target.to(rank)
optimizer.zero_grad()
output = ddp_model(data.view(data.size(0), -1))
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
dist.destroy_process_group()
if __name__ == '__main__':
world_size = 2
mp.spawn(setup, args=(world_size,), nprocs=world_size, join=True)
PS：需要说明的是，上面这个kaggle双T4GPU训练只能在命令行使用!python ./model.py的方式运行，如果直接在kaggle的Jupypter Notebook上，上面的代码运行会报错的，具体原因我没有细查，但是与Jupypter Notebook不能使用import torch.multiprocessing as mp有关，因此如果直接使用Jupypter Notebook，则推荐使用accelerate库，示例代码如下：
from accelerate import Accelerator, notebook_launcher
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
def train_mnist():
accelerator = Accelerator()
# 定义模型
model = nn.Sequential(
nn.Linear(784, 4096),
nn.ReLU(),
nn.Linear(4096, 2048),
nn.ReLU(),
nn.Linear(2048, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10)
)
# 定义优化器和损失函数
optimizer = optim.SGD(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
# 为分布式训练准备模型、优化器和损失函数
model, optimizer, loss_fn = accelerator.prepare(model, optimizer, loss_fn)
# 定义数据转换
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载 MNIST 数据集
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
# 为分布式训练准备数据加载器
train_loader = accelerator.prepare(train_loader)
# 训练循环
for epoch in range(50):
if accelerator.is_main_process:
print('当前正在进行第{}轮训练'.format(epoch))
if epoch == 1:
print('当前模型所在设备: ', next(model.parameters()).device)
pass
pass
for data, target in train_loader:
optimizer.zero_grad()
output = model(data.view(data.size(0), -1)) # 展平输入
loss = loss_fn(output, target)
accelerator.backward(loss) # 反向传播
optimizer.step()
if accelerator.is_main_process:
print("训练完成！")
# 使用 2 个进程（对应 2 个 GPU）启动训练
notebook_launcher(train_mnist, num_processes=2)
参考链接1：动画理解Pytorch 大模型分布式训练技术 DP，DDP，DeepSpeed ZeRO技术 - 哔哩哔哩
参考链接2：「分布式训练」原理讲解+ 「DDP 代码实现」修改要点 - 哔哩哔哩
参考链接3：2024大模型之AI框架之分布式训练 - 哔哩哔哩
参考链接4：一次搞懂PyTorch DDP分布式训练 - 哔哩哔哩
参考链接5：pytorch多GPU并行训练教程 - 哔哩哔哩
参考链接6：【分布式训练基础知识】01.NVIDIA硬件与软件 - 哔哩哔哩
参考链接7：Horovod的安装和使用 - Mario的文章 - 知乎
参考链接8：Accelerate库 Kaggle T4*2 单机多卡DDP数据并行训练 - LeonYi的文章 - 知乎
浙公网安备 33010602011771号