FAFormulas and AlgorithmCNNConvolutional Neural Network1. 引言在深度学习训练中GPU 的并行计算能力能大幅缩短模型训练时间。本文将基于 PyTorch 框架以 CIFAR10 图像分类任务为例详细介绍利用 GPU 训练神经网络的两种方式并对比 CPU 与 GPU 的训练耗时差异。2. 环境准备在开始之前请确保你的环境中已安装以下依赖pipinstalltorch torchvision tensorboard同时建议通过nvidia-smi命令查看当前显卡配置确认 GPU 是否可用!nvidia-smi输出示例----------------------------------------------------------------------------- | NVIDIA-SMI 471.35 Driver Version: 471.35 CUDA Version: 11.4 | |--------------------------------------------------------------------------- | GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | || | 0 NVIDIA GeForce ... WDDM | 00000000:01:00.0 On | N/A | | N/A 61C P0 47W / N/A | 2913MiB / 16384MiB | 10% Default | -----------------------------------------------------------------------------3. 利用 GPU 训练方式一3.1 核心思路GPU 训练的核心在于将以下三部分迁移到 GPU 上网络模型Model数据输入与标注损失函数Loss Function3.2 完整代码importtorchvisionimporttorchfromtorchimportnnfromtorch.utils.dataimportDataLoaderfromtorch.utils.tensorboardimportSummaryWriter# 定义网络模型classTudui(nn.Module):def__init__(self):super(Tudui,self).__init__()self.model1nn.Sequential(nn.Conv2d(3,32,5,1,2),# 输入通道3输出通道32卷积核5×5步长1填充2nn.MaxPool2d(2),nn.Conv2d(32,32,5,1,2),nn.MaxPool2d(2),nn.Conv2d(32,64,5,1,2),nn.MaxPool2d(2),nn.Flatten(),# 展平后变成 64*4*4nn.Linear(64*4*4,64),nn.Linear(64,10))defforward(self,x):xself.model1(x)returnx# 准备数据集train_datatorchvision.datasets.CIFAR10(./dataset,trainTrue,transformtorchvision.transforms.ToTensor(),downloadTrue)test_datatorchvision.datasets.CIFAR10(./dataset,trainFalse,transformtorchvision.transforms.ToTensor(),downloadTrue)train_data_sizelen(train_data)test_data_sizelen(test_data)print(训练数据集的长度{}.format(train_data_size))print(测试数据集的长度{}.format(test_data_size))# 利用 DataLoader 加载数据集train_dataloaderDataLoader(train_data,batch_size64)test_dataloaderDataLoader(test_data,batch_size64)# 创建网络模型tuduiTudui()iftorch.cuda.is_available():tuduitudui.cuda()# 网络模型转移到 GPU# 损失函数loss_fnnn.CrossEntropyLoss()iftorch.cuda.is_available():loss_fnloss_fn.cuda()# 损失函数转移到 GPU# 优化器learning0.01optimizertorch.optim.SGD(tudui.parameters(),learning)# 设置训练参数total_train_step0total_test_step0epoch10# TensorBoard 记录writerSummaryWriter(logs)foriinrange(epoch):print(-----第 {} 轮训练开始-----.format(i1))# 训练步骤tudui.train()fordataintrain_dataloader:imgs,targetsdataiftorch.cuda.is_available():imgsimgs.cuda()targetstargets.cuda()outputstudui(imgs)lossloss_fn(outputs,targets)optimizer.zero_grad()loss.backward()optimizer.step()total_train_step1iftotal_train_step%1000:print(训练次数{}Loss{}.format(total_train_step,loss.item()))writer.add_scalar(train_loss,loss.item(),total_train_step)# 测试步骤tudui.eval()total_test_loss0total_accuracy0withtorch.no_grad():fordataintest_dataloader:imgs,targetsdataiftorch.cuda.is_available():imgsimgs.cuda()targetstargets.cuda()outputstudui(imgs)lossloss_fn(outputs,targets)total_test_lossloss.item()accuracy(outputs.argmax(1)targets).sum()total_accuracyaccuracyprint(整体测试集上的Loss{}.format(total_test_loss))print(整体测试集上的正确率{}.format(total_accuracy/test_data_size))writer.add_scalar(test_loss,total_test_loss,total_test_step)writer.add_scalar(test_accuracy,total_accuracy/test_data_size,total_test_step)total_test_step1torch.save(tudui,./model/tudui_{}.pth.format(i))print(模型已保存)writer.close()3.3 运行结果训练数据集的长度50000 测试数据集的长度10000 -----第 1 轮训练开始----- 训练次数100Loss2.289992094039917 训练次数200Loss2.2927844524383545 ... 整体测试集上的Loss317.0560564994812 整体测试集上的正确率0.28700000047683716 模型已保存4. GPU 训练时间对比为了量化 GPU 加速效果我们在训练循环中加入计时逻辑统计每训练 100 次所消耗的时间。importtime start_timetime.time()foriinrange(epoch):# 训练步骤tudui.train()fordataintrain_dataloader:imgs,targetsdataiftorch.cuda.is_available():imgsimgs.cuda()targetstargets.cuda()outputstudui(imgs)lossloss_fn(outputs,targets)optimizer.zero_grad()loss.backward()optimizer.step()total_train_step1iftotal_train_step%1000:end_timetime.time()print(训练 100 次耗时{} 秒.format(end_time-start_time))print(训练次数{}Loss{}.format(total_train_step,loss.item()))writer.add_scalar(train_loss,loss.item(),total_train_step)运行结果示例训练次数100Loss2.2871038913726807 训练 100 次耗时2.1766483783721924 秒 训练次数200Loss2.2836720943450928 训练 100 次耗时3.27374267578125 秒 ...5. 利用 GPU 训练方式二5.1 核心思路方式二通过torch.device统一管理设备代码更简洁、可读性更强。当电脑有多个显卡时还可以通过cuda:0、cuda:1指定使用哪块 GPU。5.2 完整代码importtorchvisionimporttorchfromtorchimportnnfromtorch.utils.dataimportDataLoaderfromtorch.utils.tensorboardimportSummaryWriterimporttime# 定义训练的设备# device torch.device(cpu)# device torch.device(cuda) # 使用 GPU 方式一# device torch.device(cuda:0) # 使用 GPU 方式二devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)# 定义网络模型classTudui(nn.Module):def__init__(self):super(Tudui,self).__init__()self.model1nn.Sequential(nn.Conv2d(3,32,5,1,2),nn.MaxPool2d(2),nn.Conv2d(32,32,5,1,2),nn.MaxPool2d(2),nn.Conv2d(32,64,5,1,2),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64*4*4,64),nn.Linear(64,10))defforward(self,x):xself.model1(x)returnx# 准备数据集train_datatorchvision.datasets.CIFAR10(./dataset,trainTrue,transformtorchvision.transforms.ToTensor(),downloadTrue)test_datatorchvision.datasets.CIFAR10(./dataset,trainFalse,transformtorchvision.transforms.ToTensor(),downloadTrue)train_data_sizelen(train_data)test_data_sizelen(test_data)print(训练数据集的长度{}.format(train_data_size))print(测试数据集的长度{}.format(test_data_size))train_dataloaderDataLoader(train_data,batch_size64)test_dataloaderDataLoader(test_data,batch_size64)# 创建网络模型并迁移到指定设备tuduiTudui()tuduitudui.to(device)# 损失函数loss_fnnn.CrossEntropyLoss()loss_fnloss_fn.to(device)# 优化器learning0.01optimizertorch.optim.SGD(tudui.parameters(),learning)total_train_step0total_test_step0epoch10writerSummaryWriter(logs)start_timetime.time()foriinrange(epoch):print(-----第 {} 轮训练开始-----.format(i1))tudui.train()fordataintrain_dataloader:imgs,targetsdata imgsimgs.to(device)targetstargets.to(device)outputstudui(imgs)lossloss_fn(outputs,targets)optimizer.zero_grad()loss.backward()optimizer.step()total_train_step1iftotal_train_step%1000:end_timetime.time()print(训练 100 次耗时{} 秒.format(end_time-start_time))print(训练次数{}Loss{}.format(total_train_step,loss.item()))writer.add_scalar(train_loss,loss.item(),total_train_step)tudui.eval()total_test_loss0total_accuracy0withtorch.no_grad():fordataintest_dataloader:imgs,targetsdata imgsimgs.to(device)targetstargets.to(device)outputstudui(imgs)lossloss_fn(outputs,targets)total_test_lossloss.item()accuracy(outputs.argmax(1)targets).sum()total_accuracyaccuracyprint(整体测试集上的Loss{}.format(total_test_loss))print(整体测试集上的正确率{}.format(total_accuracy/test_data_size))writer.add_scalar(test_loss,total_test_loss,total_test_step)writer.add_scalar(test_accuracy,total_accuracy/test_data_size,total_test_step)total_test_step1torch.save(tudui,./model/tudui_{}.pth.format(i))print(模型已保存)writer.close()5.3 运行结果训练数据集的长度50000 测试数据集的长度10000 -----第 1 轮训练开始----- 训练 100 次耗时1.1190404891967773 秒 训练次数100Loss2.2926671504974365 ... 整体测试集上的Loss195.62357383966446 整体测试集上的正确率0.5604999661445618 模型已保存6. 运行 Terminal 语句在 Jupyter Notebook 或 Colab 中可以通过在代码块中语句前添加感叹号!来执行终端命令。例如查看显卡配置!nvidia-smi7. 总结本文介绍了 PyTorch 利用 GPU 训练神经网络的两种方式方式特点适用场景方式一使用cuda()方法逐个迁移模型、损失函数和数据简单直观适合初学者方式二使用torch.device统一管理设备代码简洁支持多 GPU 指定推荐使用通过 GPU 加速训练效率显著提升。建议在实际项目中优先采用方式二便于后续扩展和代码维护。