使用GPU训练模型
本文代码详解参考:
模型训练基础流程-CSDN博客
目录
为什么要用GPU训练模型
什么是CUDA
利用GPU训练—方式一(.cuda())
利用GPU训练—方式二 (.to())
Google Colaboratory
为什么要用GPU训练模型
用 GPU 训练模型的核心原因是GPU 的硬件架构和计算特性,能完美适配深度学习模型的计算需求,大幅提升训练效率
一、GPU 的核心优势:并行计算能力极强
深度学习模型的训练过程(如神经网络的前向传播、反向传播)本质上是大量重复的矩阵运算和向量运算(例如卷积层的卷积操作、全连接层的矩阵乘法)。
-
CPU 的设计侧重 “低延迟、串行计算”,核心数量少(通常 4-32 核),适合处理逻辑复杂、步骤连贯的任务(如系统调度、单线程指令),但面对成百上千万次的并行重复计算时效率极低。
-
GPU 的设计则侧重 “高吞吐、并行计算”,核心数量极多(主流 GPU 有数千个计算核心,如 RTX 4090 有 16384 个 CUDA 核心),这些核心可以同时处理大量相似的简单计算(比如同时对矩阵中不同位置的元素做乘法)。
例如:一个 1000×1000 的矩阵乘法,CPU 可能需要逐个元素循环计算,而 GPU 可以同时启动上万次运算,将计算时间从 “小时级” 压缩到 “分钟级” 甚至 “秒级”。
二、适配深度学习的 “计算密集型” 需求
深度学习训练的两大核心特点,恰好被 GPU 针对性解决:
-
计算量极大:模型参数量从百万级(简单 CNN)到千亿级(大语言模型),每次迭代需要对所有参数计算梯度,涉及的运算次数按 “亿” 甚至 “万亿” 计。GPU 的并行核心能同时分摊这些计算,避免 CPU “逐个处理” 的低效。
-
数据吞吐量高:训练时需要频繁读取批量数据(如图像、文本特征),并在模型各层间传递。GPU 自带大容量高带宽显存(如 16GB-80GB GDDR6/HOF),配合专门的存储控制器,能快速读写数据,避免 “数据等待计算” 的瓶颈(而 CPU 的内存带宽通常仅为 GPU 的 1/10 左右)。
三、实际效果:训练效率提升数十倍甚至上百倍
-
对于简单模型(如小型 CNN):GPU 训练速度通常是 CPU 的 10-30 倍。
-
对于大型模型(如 Transformer、大语言模型):CPU 可能需要数周甚至数月才能完成训练,而 GPU(尤其是多 GPU 集群)可压缩到几天甚至几小时,且能支持更大的批量和更复杂的模型结构(否则 CPU 会因内存或速度限制无法运行)。
用后述代码展示速度差异:
利用GPU训练的时间:
利用CPU训练的时间:
CPU训练百次需要2s左右 而GPU只要0.5s ! ! !
可以使用如图命令查看自己设备的GPU信息
若报错则是驱动没有正确安装,需要去英伟达官网下载
什么是CUDA
CUDA(Compute Unified Device Architecture)是英伟达(NVIDIA)推出的一种并行计算平台和编程模型,旨在利用 NVIDIA GPU 的并行计算能力,加速计算密集型任务。
核心概念
-
并行计算平台:CUDA 为开发者提供了一个环境,使得 GPU 可以作为一个高度并行的计算设备,和 CPU 协同工作。在传统计算中,CPU 处理任务是串行的,而 CUDA 允许将大量可以并行执行的计算任务分配到 GPU 上,极大地提高计算效率。
-
编程模型:CUDA 定义了一套编程接口和语法规则,允许开发者使用类 C 语言(CUDA C/C++)编写并行计算程序,来充分发挥 GPU 的大规模并行计算能力。此外,像 PyTorch、TensorFlow 等深度学习框架也对 CUDA 进行了封装,使得深度学习开发者无需深入了解底层编程细节,就能轻松利用 GPU 加速模型训练。
工作原理
-
任务拆分:在 CUDA 编程中,开发者需要将计算任务分解成大量可以并行执行的小任务。例如,在深度学习的矩阵乘法中,矩阵的不同元素运算可以并行执行,CUDA 可以将这些运算分配到 GPU 的众多计算核心上。
-
线程管理:CUDA 将 GPU 的计算资源组织成线程层次结构,包括线程块(block)和线程网格(grid)。每个线程执行相同的内核函数(kernel function),但处理不同的数据,通过线程 ID 来区分和操作不同的数据。
-
协同计算:CPU 负责处理逻辑性强、计算量小的任务,如程序的流程控制、数据的预处理和后处理等;GPU 则专注于执行高度并行的计算任务,两者相互协作,共同完成复杂的计算任务。
应用场景
-
深度学习:是 CUDA 最广泛的应用领域之一。在训练神经网络模型(如卷积神经网络 CNN、循环神经网络 RNN 及其变体 LSTM、GRU,还有 Transformer 等)时,大量的矩阵运算(如前向传播、反向传播)可以利用 CUDA 在 GPU 上并行加速,大幅缩短训练时间。
-
科学计算:在物理模拟、流体动力学计算、分子动力学模拟等科学领域,存在大量的数值计算,CUDA 能够显著提升这些计算的效率,加速科研进程。
-
图形渲染:虽然 GPU 最初是为图形渲染设计的,但 CUDA 进一步拓展了其应用。在实时渲染、光线追踪等图形处理任务中,CUDA 可以加速图形算法的执行,提升渲染质量和速度。
利用GPU训练—方式一(.cuda())
把如图所示的三个 加上cuda方法即可
运行后GPU确实跑起来了:
代码:
import torchimport torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import timetrain_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,transform=torchvision.transforms.ToTensor(),download=True)test_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,transform=torchvision.transforms.ToTensor(),download=False)# 数据集长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))# 利用DataLoader来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 创建网络模型
class MyModule(nn.Module):def __init__(self):super().__init__()self.module = nn.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))def forward(self, x):x = self.module(x)return xmodel = MyModule()
# 网络模型转移到GPU
if torch.cuda.is_available():model = model.cuda()# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 损失函数转移到GPU
if torch.cuda.is_available():loss_fn = loss_fn.cuda()
# 优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)# 设置训练网络的一些参数
# 训练次数
total_train_step = 0
# 测试次数
total_test_step = 0
# 训练轮数
epoch = 10# 添加tensorboard
writer = SummaryWriter("logs_test16")# 计时比较CPU和GPU的训练速度差异
# 开始时间
start_time=time.time()
for i in range(epoch):print("---------第{}轮训练开始---------".format(i + 1))# 训练model.train()for data in train_dataloader:imgs, targets = data# 训练数据转移到GPUif torch.cuda.is_available():imgs = imgs.cuda()targets = targets.cuda()outputs = model(imgs)loss = loss_fn(outputs, targets)# 优化器优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_train_step = total_train_step + 1if total_train_step % 100 == 0:# 到当前的训练总时间end_time=time.time()print("训练所花时间:{}秒".format(end_time-start_time))print("训练次数:{},Loss:{}".format(total_train_step, loss.item()))writer.add_scalar("train_loss", loss.item(), total_train_step)# 评估模型训练效果 - 跑一个测试数据查看正确率# 测试过程中,参数不需要调整了,临时把梯度去除model.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_dataloader:imgs, targets = data# 测试数据转移到GPUif torch.cuda.is_available():imgs = imgs.cuda()targets = targets.cuda()outputs = model(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()accuracy = (outputs.argmax(1) == targets).sum()total_accuracy = total_accuracy + accuracy # 预测对的总数量print("整体测试集上的loss:{}", format(total_test_loss))print("整体测试集上的准确率:{}", format(total_accuracy / test_data_size)) # 准确率writer.add_scalar("test_sum_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)total_test_step = total_test_step + 1# 保存每一次训练的结果torch.save(model, "model_{}.pth".format(i))print("模型已保存")writer.close()
import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import timetrain_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,
transform=torchvision.transforms.ToTensor(),
download=True)test_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,
transform=torchvision.transforms.ToTensor(),
download=False)# 数据集长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))# 利用DataLoader来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 创建网络模型
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.module = nn.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)
)def forward(self, x):
x = self.module(x)
return xmodel = MyModule()
# 网络模型转移到GPU
if torch.cuda.is_available():
model = model.cuda()# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 损失函数转移到GPU
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
# 优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)# 设置训练网络的一些参数
# 训练次数
total_train_step = 0
# 测试次数
total_test_step = 0
# 训练轮数
epoch = 10# 添加tensorboard
writer = SummaryWriter("logs_test16")# 计时比较CPU和GPU的训练速度差异
# 开始时间
start_time=time.time()
for i in range(epoch):
print("---------第{}轮训练开始---------".format(i + 1))# 训练
model.train()
for data in train_dataloader:
imgs, targets = data# 训练数据转移到GPU
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = model(imgs)
loss = loss_fn(outputs, targets)# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
if total_train_step % 100 == 0:
# 到当前的训练总时间
end_time=time.time()
print("训练所花时间:{}秒".format(end_time-start_time))
print("训练次数:{},Loss:{}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 评估模型训练效果 - 跑一个测试数据查看正确率# 测试过程中,参数不需要调整了,临时把梯度去除
model.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data# 测试数据转移到GPU
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()outputs = model(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss.item()
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy = total_accuracy + accuracy # 预测对的总数量print("整体测试集上的loss:{}", format(total_test_loss))
print("整体测试集上的准确率:{}", format(total_accuracy / test_data_size)) # 准确率
writer.add_scalar("test_sum_loss", total_test_loss, total_test_step)
writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
total_test_step = total_test_step + 1# 保存每一次训练的结果
torch.save(model, "model_{}.pth".format(i))
print("模型已保存")writer.close()
利用GPU训练—方式二 (.to())
更常用和推荐
import torchimport torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time# 定义训练的设备
device1 = torch.device("cpu")
device2 = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device2 = torch.device("cuda") 在确保有GPU的情况下也可以train_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,transform=torchvision.transforms.ToTensor(),download=True)test_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,transform=torchvision.transforms.ToTensor(),download=False)# 数据集长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))# 利用DataLoader来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 创建网络模型
class MyModule(nn.Module):def __init__(self):super().__init__()self.module = nn.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))def forward(self, x):x = self.module(x)return xmodel = MyModule()
# 网络模型转移到GPU
model = model.to(device2)# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 损失函数转移到GPU
loss_fn = loss_fn.to(device2)
# 优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)# 设置训练网络的一些参数
# 训练次数
total_train_step = 0
# 测试次数
total_test_step = 0
# 训练轮数
epoch = 10# 添加tensorboard
writer = SummaryWriter("logs_test17")# 计时比较CPU和GPU的训练速度差异
# 开始时间
start_time = time.time()
for i in range(epoch):print("---------第{}轮训练开始---------".format(i + 1))# 训练model.train()for data in train_dataloader:imgs, targets = data# 训练数据转移到GPUimgs = imgs.to(device2)targets = targets.to(device2)outputs = model(imgs)loss = loss_fn(outputs, targets)# 优化器优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_train_step = total_train_step + 1if total_train_step % 100 == 0:# 到当前的训练总时间end_time = time.time()print("训练所花时间:{}秒".format(end_time - start_time))print("训练次数:{},Loss:{}".format(total_train_step, loss.item()))writer.add_scalar("train_loss", loss.item(), total_train_step)# 评估模型训练效果 - 跑一个测试数据查看正确率# 测试过程中,参数不需要调整了,临时把梯度去除model.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_dataloader:imgs, targets = data# 测试数据转移到GPUimgs = imgs.to(device2)targets = targets.to(device2)outputs = model(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()accuracy = (outputs.argmax(1) == targets).sum()total_accuracy = total_accuracy + accuracy # 预测对的总数量print("整体测试集上的loss:{}", format(total_test_loss))print("整体测试集上的准确率:{}", format(total_accuracy / test_data_size)) # 准确率writer.add_scalar("test_sum_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)total_test_step = total_test_step + 1# 保存每一次训练的结果torch.save(model, "model_{}.pth".format(i))print("模型已保存")writer.close()
import torchimport torchvision from torch import nn from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter import time# 定义训练的设备 device1 = torch.device("cpu") device2 = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,transform=torchvision.transforms.ToTensor(),download=True)test_data = torchvision.datasets.CIFAR10(root="D:\\Python\\learn_pytorch\\torchvision_dataset", train=True,transform=torchvision.transforms.ToTensor(),download=False)# 数据集长度 train_data_size = len(train_data) test_data_size = len(test_data) print("训练数据集的长度为:{}".format(train_data_size)) print("测试数据集的长度为:{}".format(test_data_size))# 利用DataLoader来加载数据集 train_dataloader = DataLoader(train_data, batch_size=64) test_dataloader = DataLoader(test_data, batch_size=64)# 创建网络模型 class MyModule(nn.Module):def __init__(self):super().__init__()self.module = nn.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))def forward(self, x):x = self.module(x)return xmodel = MyModule() # 网络模型转移到GPU model = model.to(device2)# 损失函数 loss_fn = nn.CrossEntropyLoss() # 损失函数转移到GPU loss_fn = loss_fn.to(device2) # 优化器 learning_rate = 1e-2 optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)# 设置训练网络的一些参数 # 训练次数 total_train_step = 0 # 测试次数 total_test_step = 0 # 训练轮数 epoch = 10# 添加tensorboard writer = SummaryWriter("logs_test17")# 计时比较CPU和GPU的训练速度差异 # 开始时间 start_time = time.time() for i in range(epoch):print("---------第{}轮训练开始---------".format(i + 1))# 训练model.train()for data in train_dataloader:imgs, targets = data# 训练数据转移到GPUimgs = imgs.to(device2)targets = targets.to(device2)outputs = model(imgs)loss = loss_fn(outputs, targets)# 优化器优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_train_step = total_train_step + 1if total_train_step % 100 == 0:# 到当前的训练总时间end_time = time.time()print("训练所花时间:{}秒".format(end_time - start_time))print("训练次数:{},Loss:{}".format(total_train_step, loss.item()))writer.add_scalar("train_loss", loss.item(), total_train_step)# 评估模型训练效果 - 跑一个测试数据查看正确率# 测试过程中,参数不需要调整了,临时把梯度去除model.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_dataloader:imgs, targets = data # 测试数据转移到GPUimgs = imgs.to(device2)targets = targets.to(device2)outputs = model(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()accuracy = (outputs.argmax(1) == targets).sum()total_accuracy = total_accuracy + accuracy # 预测对的总数量print("整体测试集上的loss:{}", format(total_test_loss))print("整体测试集上的准确率:{}", format(total_accuracy / test_data_size)) # 准确率writer.add_scalar("test_sum_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)total_test_step = total_test_step + 1# 保存每一次训练的结果torch.save(model, "model_{}.pth".format(i))print("模型已保存")writer.close()
Google Colaboratory
Google Colab(Colaboratory)是谷歌开发的基于云端的交互式笔记本环境。
- 核心功能:无需本地配置环境,可直接在浏览器中编写和运行 Python 代码,支持深度学习、数据分析等任务。
- 硬件支持:免费提供 CPU、GPU(如 Tesla K80、T4)甚至 TPU(张量处理单元)资源,方便运行大型模型(如训练神经网络)。
- 优势:
- 与 Google Drive 无缝集成,可直接读取 / 保存云端文件;
- 支持实时协作(多人共同编辑同一笔记本);
- 预装了 PyTorch、TensorFlow、NumPy 等主流库,开箱即用。
- 适用场景:快速原型开发、学习深度学习、资源有限时的模型训练等。
简单说,它是一个 “云端免费 GPU 编程工具”,对初学者和需要临时算力的开发者非常友好。
前提:能使用谷歌 (-_-)
在上方导航栏 修改-笔记本设置 里可以启用GPU/TPU
查看硬件参数
Tesla T4 是一款面向数据中心的推理加速 GPU,基于 Turing 架构,具备不错的计算能力,支持多种计算精度,在深度学习推理任务中有出色表现。
用 Colaboratory 跑上述代码