Pytorch 的实际应用学习笔记

一. 模型的下载

weights为false时则为没有提前经过训练的模型，为true时则经过了提前训练

vgg16_false = torchvision.models.vgg16(weights=False)
vgg16_true = torchvision.models.vgg16(weights=True)

打印

二. 模型的修改

（1）添加操作

分为两种，一种是在classifier的外部添加，一种是在内部添加

外部添加，例如添加了一个线性层

vgg16_true.add_module("add_linear", nn.Linear(1000, 10))

打印，最下方添加了线性层

内部添加

vgg16_true.classifier.add_module("add_linear", nn.Linear(1000, 10))

打印，在classifier里面添加

（2）修改操作

例如，我修改索引为6的操作

vgg16_false.classifier[6] = nn.Linear(4096, 5)

打印

三. 模型的保存与读取

在PyTorch中，可以使用torch.save函数来保存模型的状态字典或整个模型。同时，可以使用torch.load函数来加载保存的模型。

1. 保存模型的状态字典：

# 保存模型的状态字典
torch.save(model.state_dict(), 'model.pth')# 加载模型的状态字典
model.load_state_dict(torch.load('model.pth'))

2. 保存整个模型：

# 保存整个模型
torch.save(model, 'model.pth')# 加载整个模型
model = torch.load('model.pth')

需要注意的是，如果要加载模型，需要确保模型的定义和保存时一致。如果要加载模型到GPU上，需要在torch.load函数中传入map_location参数来指定加载到哪个设备上。

四. 训练套路实例

训练流程

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterfrom ch2.model import Yktrain_set = torchvision.datasets.CIFAR10("./data", transform=torchvision.transforms.ToTensor(), train=True,download=True)
test_set = torchvision.datasets.CIFAR10("./data", transform=torchvision.transforms.ToTensor(), train=False,download=True)
train_loader = DataLoader(dataset=train_set, batch_size=64, drop_last=True)
test_loader = DataLoader(dataset=test_set, batch_size=64, drop_last=True)loss_fn = nn.CrossEntropyLoss()learning_rate = 1e-2
yk = Yk()
opt = torch.optim.SGD(yk.model1.parameters(), learning_rate)total_train_step = 0
total_test_step = 0
epoch = 10
writer = SummaryWriter("../logs")
for i in range(epoch):print("————第{}次训练开始————".format(i))yk.train()for data in train_loader:images, targets = dataoutput = yk(images)loss = loss_fn(output, targets)opt.zero_grad()loss.backward()opt.step()total_train_step += 1if total_train_step % 100 == 0:print("训练次数:{}, Loss:{}".format(total_train_step, loss))writer.add_scalar("train_loss", loss.item(), total_train_step)# 测试步骤开始yk.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_loader:images, targets = dataoutput = yk(images)loss = loss_fn(output, targets)total_test_loss = total_test_loss + lossaccuracy = (output.argmax(1) == targets).sum()total_accuracy=total_accuracy+accuracyprint("整体测试集上的loss:{}".format(total_test_step))print("整体测试集上的正确率{}".format(total_accuracy/len(test_set)))total_test_step += 1writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)writer.add_scalar("test_accuracy", total_accuracy.item(), total_test_step)torch.save(yk, "yk_{}".format(i))print("模型已保存")writer.close()

训练模型

from torch import nnclass Yk(nn.Module):def __init__(self):super(Yk, self).__init__()self.model1 = nn.Sequential(nn.Conv2d(3, 32, (5, 5), padding=2),nn.MaxPool2d(2),nn.Conv2d(32, 32, (5, 5), padding=2),nn.MaxPool2d(2),nn.Conv2d(32, 64, (5, 5), padding=2),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(4 * 4 * 64, 64),nn.Linear(64, 10))def forward(self, x):x = self.model1(x)return xif __name__ == '__main__':yk = Yk()

打印

五. 使用GPU训练

1. 使用cuda

原本代码

import timeimport torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterfrom ch2.model import Yktrain_set = torchvision.datasets.CIFAR10("./data", transform=torchvision.transforms.ToTensor(), train=True,download=True)
test_set = torchvision.datasets.CIFAR10("./data", transform=torchvision.transforms.ToTensor(), train=False,download=True)
train_loader = DataLoader(dataset=train_set, batch_size=64, drop_last=True)
test_loader = DataLoader(dataset=test_set, batch_size=64, drop_last=True)start_time = time.time()loss_fn = nn.CrossEntropyLoss()
# loss_fn = loss_fn.cuda()
learning_rate = 1e-2
yk = Yk()
# yk = yk.cuda()
opt = torch.optim.SGD(yk.model1.parameters(), learning_rate)total_train_step = 0
total_test_step = 0
epoch = 10
writer = SummaryWriter("../logs")
for i in range(epoch):print("————第{}次训练开始————".format(i))yk.train()for data in train_loader:images, targets = data# images = images.cuda()# targets = targets.cuda()output = yk(images)loss = loss_fn(output, targets)opt.zero_grad()loss.backward()opt.step()total_train_step += 1if total_train_step % 100 == 0:end_time = time.time()print(end_time-start_time)print("训练次数:{}, Loss:{}".format(total_train_step, loss))writer.add_scalar("train_loss", loss.item(), total_train_step)# 测试步骤开始yk.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_loader:images, targets = data# images = images.cuda()# targets = targets.cuda()output = yk(images)loss = loss_fn(output, targets)total_test_loss = total_test_loss + lossaccuracy = (output.argmax(1) == targets).sum()total_accuracy = total_accuracy + accuracyprint("整体测试集上的loss:{}".format(total_test_step))print("整体测试集上的正确率{}".format(total_accuracy / len(test_set)))total_test_step += 1writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)writer.add_scalar("test_accuracy", total_accuracy.item(), total_test_step)torch.save(yk, "yk_{}".format(i))print("模型已保存")writer.close()

间隔时间如下

使用cuda，用gpu后，间隔时间明显极大缩短了

2. 定义设备

在里面定义使用cpu，还是cuda，有多个gpu，可以选用第几个（我选用的第一个）

device = torch.device("cuda:0")

原本需要使用*.cuda的地方，修改为*.to(device)：

yk = yk.to(device)

全部代码如下：

import timeimport torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterfrom ch2.model import Ykdevice = torch.device("cuda:0")train_set = torchvision.datasets.CIFAR10("./data", transform=torchvision.transforms.ToTensor(), train=True,download=True)
test_set = torchvision.datasets.CIFAR10("./data", transform=torchvision.transforms.ToTensor(), train=False,download=True)
train_loader = DataLoader(dataset=train_set, batch_size=64, drop_last=True)
test_loader = DataLoader(dataset=test_set, batch_size=64, drop_last=True)start_time = time.time()loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)
learning_rate = 1e-2
yk = Yk()
yk = yk.to(device)
opt = torch.optim.SGD(yk.model1.parameters(), learning_rate)total_train_step = 0
total_test_step = 0
epoch = 10
writer = SummaryWriter("../logs")
for i in range(epoch):print("————第{}次训练开始————".format(i))yk.train()for data in train_loader:images, targets = dataimages = images.to(device)targets = targets.to(device)output = yk(images)loss = loss_fn(output, targets)opt.zero_grad()loss.backward()opt.step()total_train_step += 1if total_train_step % 100 == 0:end_time = time.time()print(end_time - start_time)print("训练次数:{}, Loss:{}".format(total_train_step, loss))writer.add_scalar("train_loss", loss.item(), total_train_step)# 测试步骤开始yk.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_loader:images, targets = dataimages = images.to(device)targets = targets.to(device)output = yk(images)loss = loss_fn(output, targets)total_test_loss = total_test_loss + lossaccuracy = (output.argmax(1) == targets).sum()total_accuracy = total_accuracy + accuracyprint("整体测试集上的loss:{}".format(total_test_step))print("整体测试集上的正确率{}".format(total_accuracy / len(test_set)))total_test_step += 1writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)writer.add_scalar("test_accuracy", total_accuracy.item(), total_test_step)torch.save(yk, "yk_{}".format(i))print("模型已保存")writer.close()

运行