利用GPU训练
有两种GPU训练的方式
方式1：
网络模型
损失函数
数据（输入，标注）
.cuda()
if torch.cuda.is_available():
pongber = pongber.cuda()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
加上if torch.cuda.is_available():这个判断，这个程序无论在CPU上还是在GPU上都可以跑，优先在GPU上跑
import torch
import torchvision
from numpy.ma.core import argmax
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
# length 长度 查看数据集有多少
train_data_size = len(train_data)
test_data_size = len(test_data)
# 如果train_data.size=10，训练数据集的长度为：10
print("训练数据集的长度为：{}".format(train_data_size)) # 字符串格式化
print("测试数据集的长度为：{}".format(test_data_size)) # 字符串格式化
print("训练数据集的个数为%d，测试数据集的个数为%d" %(train_data_size, test_data_size))
# 利用Dataloader来加载数据集
train_dataLoader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 搭建神经网络
class Pongber(nn.Module):
def __init__(self):
super(Pongber, self).__init__()
self.model = 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(1024, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 创建网络模型
pongber = Pongber()
if torch.cuda.is_available():
pongber = pongber.cuda()
# 损失函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
# 优化器
learning_rate = 0.01
# 1e-2 1e-2 = 1×(10)^(-2) = 1/100 = 0.01
# learning_rate = 1e-2
optimizer = torch.optim.SGD(pongber.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter(log_dir="logs_train")
# epoch = 10，i 从0跑到9
for i in range(epoch):
print("------第{}轮训练开始------".format(i+1))
# 训练步骤开始
pongber.train(mode=True)
for data in train_dataLoader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
# print("训练次数：{}， Loss：{}".format(total_train_step, loss))
if total_train_step % 100 == 0: # 每100次打印记录，避免一些无用的信息
print("训练次数：{}， Loss：{}".format(total_train_step, loss.item())) # 加上item就会把它转换成一个真实的数字
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤开始
pongber.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad(): # 没有梯度，保证不进行调优
for data in test_dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
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_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(pongber, "P28_train_gpu_1/pongber_{}.pth".format(i+1))
# torch.save(pongber.state_dict(), "pongber_{}.pth".format(i+1))
print("第{}轮训练模型已经保存".format(i+1))
writer.close()
比较一下GPU和CPU训练时间差异
start_time = time.time()
end_time = time.time()
print(end_time - start_time)
训练过程中，在终端输入nvidia-smi，会输出一些GPU的信息
GPU版本
import torch
import torchvision
from numpy.ma.core import argmax
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
# length 长度 查看数据集有多少
train_data_size = len(train_data)
test_data_size = len(test_data)
# 如果train_data.size=10，训练数据集的长度为：10
print("训练数据集的长度为：{}".format(train_data_size)) # 字符串格式化
print("测试数据集的长度为：{}".format(test_data_size)) # 字符串格式化
print("训练数据集的个数为%d，测试数据集的个数为%d" %(train_data_size, test_data_size))
# 利用Dataloader来加载数据集
train_dataLoader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 搭建神经网络
class Pongber(nn.Module):
def __init__(self):
super(Pongber, self).__init__()
self.model = 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(1024, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 创建网络模型
pongber = Pongber()
if torch.cuda.is_available():
pongber = pongber.cuda()
# 损失函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
# 优化器
learning_rate = 0.01
# 1e-2 1e-2 = 1×(10)^(-2) = 1/100 = 0.01
# learning_rate = 1e-2
optimizer = torch.optim.SGD(pongber.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter(log_dir="logs_train")
start_time = time.time()
# epoch = 10，i 从0跑到9
for i in range(epoch):
print("------第{}轮训练开始------".format(i+1))
# 训练步骤开始
pongber.train(mode=True)
for data in train_dataLoader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
# print("训练次数：{}， Loss：{}".format(total_train_step, loss))
if total_train_step % 100 == 0: # 每100次打印记录，避免一些无用的信息
end_time = time.time()
print(end_time - start_time)
print("训练次数：{}， Loss：{}".format(total_train_step, loss.item())) # 加上item就会把它转换成一个真实的数字
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤开始
pongber.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad(): # 没有梯度，保证不进行调优
for data in test_dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
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_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(pongber, "P28_train_gpu_1/pongber_{}.pth".format(i+1))
# torch.save(pongber.state_dict(), "pongber_{}.pth".format(i+1))
print("第{}轮训练模型已经保存".format(i+1))
writer.close()
CPU版本
import torch
import torchvision
from numpy.ma.core import argmax
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
# length 长度 查看数据集有多少
train_data_size = len(train_data)
test_data_size = len(test_data)
# 如果train_data.size=10，训练数据集的长度为：10
print("训练数据集的长度为：{}".format(train_data_size)) # 字符串格式化
print("测试数据集的长度为：{}".format(test_data_size)) # 字符串格式化
print("训练数据集的个数为%d，测试数据集的个数为%d" %(train_data_size, test_data_size))
# 利用Dataloader来加载数据集
train_dataLoader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 搭建神经网络
class Pongber(nn.Module):
def __init__(self):
super(Pongber, self).__init__()
self.model = 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(1024, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 创建网络模型
pongber = Pongber()
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 优化器
learning_rate = 0.01
# 1e-2 1e-2 = 1×(10)^(-2) = 1/100 = 0.01
# learning_rate = 1e-2
optimizer = torch.optim.SGD(pongber.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter(log_dir="logs_train")
start_time = time.time()
# epoch = 10，i 从0跑到9
for i in range(epoch):
print("------第{}轮训练开始------".format(i+1))
# 训练步骤开始
pongber.train(mode=True)
for data in train_dataLoader:
imgs, targets = data
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
# print("训练次数：{}， Loss：{}".format(total_train_step, loss))
if total_train_step % 100 == 0: # 每100次打印记录，避免一些无用的信息
end_time = time.time()
print(end_time - start_time)
print("训练次数：{}， Loss：{}".format(total_train_step, loss.item())) # 加上item就会把它转换成一个真实的数字
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤开始
pongber.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad(): # 没有梯度，保证不进行调优
for data in test_dataloader:
imgs, targets = data
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
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_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(pongber, "P28_train_cpu_2/pongber_{}.pth".format(i+1))
# torch.save(pongber.state_dict(), "pongber_{}.pth".format(i+1))
print("第{}轮训练模型已经保存".format(i+1))
writer.close()
新建笔记本 -> 修改 -> 笔记本设置 -> T4 GPU ->保存
即可正常使用GPU训练
会发现它速度很快
我们想看一下它的配置，因为以上运行的都是Python语言，那如何运行在terminal上的运行的呢？
加一个感叹号，!nvidia-smi
Wed Jun 4 06:52:29 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: 12.4 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |
| N/A 59C P0 29W / 70W | 244MiB / 15360MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+
方式2：
.to(device)
device = torch.device("cpu")
torch.device("cuda")
torch.device("cuda:0") # 电脑上有一张显卡的话和上边的写法没有区别
torch.device("cuda:1") # 电脑上有多个显卡，指定第2张显卡
import torch
import torchvision
from numpy.ma.core import argmax
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
# 定义训练的设备
# device = torch.device("cpu")
device = torch.device("cuda:0")
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="P27_train/dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
# length 长度 查看数据集有多少
train_data_size = len(train_data)
test_data_size = len(test_data)
# 如果train_data.size=10，训练数据集的长度为：10
print("训练数据集的长度为：{}".format(train_data_size)) # 字符串格式化
print("测试数据集的长度为：{}".format(test_data_size)) # 字符串格式化
print("训练数据集的个数为%d，测试数据集的个数为%d" %(train_data_size, test_data_size))
# 利用Dataloader来加载数据集
train_dataLoader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 搭建神经网络
class Pongber(nn.Module):
def __init__(self):
super(Pongber, self).__init__()
self.model = 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(1024, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 创建网络模型
pongber = Pongber()
pongber = pongber.to(device)
# 损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)
# 优化器
learning_rate = 0.01
# 1e-2 1e-2 = 1×(10)^(-2) = 1/100 = 0.01
# learning_rate = 1e-2
optimizer = torch.optim.SGD(pongber.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练的轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter(log_dir="logs_train")
start_time = time.time()
# epoch = 10，i 从0跑到9
for i in range(epoch):
print("------第{}轮训练开始------".format(i+1))
# 训练步骤开始
pongber.train(mode=True)
for data in train_dataLoader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
# print("训练次数：{}， Loss：{}".format(total_train_step, loss))
if total_train_step % 100 == 0: # 每100次打印记录，避免一些无用的信息
end_time = time.time()
print(end_time - start_time)
print("训练次数：{}， Loss：{}".format(total_train_step, loss.item())) # 加上item就会把它转换成一个真实的数字
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤开始
pongber.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad(): # 没有梯度，保证不进行调优
for data in test_dataloader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = pongber(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
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_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(pongber, "P28_train_gpu_1/pongber_{}.pth".format(i+1))
# torch.save(pongber.state_dict(), "pongber_{}.pth".format(i+1))
print("第{}轮训练模型已经保存".format(i+1))
writer.close()
这种方式是我们更加常用的方式，还有一些细节
第一个细节就是，对于这种网络模型不需要给它另外赋值
pongber = pongber.to(device)可以写成pongber.to(device)
pongber = pongber.cuda()也是可以写成pongber.cuda()
包括损失函数，只有数据、图片、标注需要转移之后再重新赋值给这个变量，但是为了方便记忆，赋值过去没有什么问题，看到别人代码这样写，不用去怀疑它有错
另外一个细节就是device = torch.device("cuda")与device = torch.device("cuda:0")对于单显卡没有区别
另外一种device = torch.device("cuda" if torch.cuda.is_available() else "cpu")，这个是一个语法的简写，cuda可用就使用gpu，不可用就使用cpu
原始资料地址：
利用GPU训练（一）
利用GPU训练（二）
如有侵权联系删除 仅供学习交流使用