PyTorch入门学习（十五）：现有网络模型的使用及修改

目录

一、使用现有的深度学习模型

二、修改深度学习模型结构

三、直接修改深度学习模型的结构

---

一、使用现有的深度学习模型

在深度学习中，许多流行的深度学习框架（如PyTorch和TensorFlow）提供了一系列现成的深度学习模型，这些模型经过在大规模数据集上的训练，并可以用于各种计算机视觉任务。在PyTorch中，使用torchvision.models模块可以轻松地加载这些模型。

下面是如何使用PyTorch中的VGG16模型的示例代码：

import torchvision
from torchvision.models import vgg16

# 加载VGG16模型
model = vgg16(pretrained=True)

# 输出模型结构
print(model)

在上面的代码中，首先导入了PyTorch的torchvision模块，然后使用vgg16(pretrained=True)来加载预训练的VGG16模型。通过print(model)，我们可以查看模型的结构。

二、修改深度学习模型结构

有时候，需要修改现有模型的结构以满足特定的任务需求。例如，如果希望将预训练的VGG16模型用于CIFAR-10数据集，我们需要更改模型的输出层，以便适应CIFAR-10的10个类别。下面是如何添加自定义输出层的示例代码：

import torch.nn as nn

# 加载VGG16模型
model = vgg16(pretrained=True)

# 添加自定义输出层，将模型的输出从1000类别转换为10类别
model.classifier.add_module('add_linear', nn.Linear(1000, 10))

# 输出修改后的模型结构
print(model)

在这个示例中，首先加载了预训练的VGG16模型，然后使用classifier.add_module方法添加了一个新的全连接层，将模型的输出从1000类别转换为10类别。

三、直接修改深度学习模型的结构

如果需要直接修改深度学习模型的结构，例如更改某一层的结构，也可以轻松实现。以下是一个示例，演示如何将VGG16模型的倒数第二个全连接层的输出维度从4096更改为10：

import torch.nn as nn

# 加载VGG16模型
model = vgg16(pretrained=True)

# 直接修改倒数第二个全连接层的输出维度
model.classifier[6] = nn.Linear(4096, 10)

# 输出修改后的模型结构
print(model)

完整代码如下：

import torch
import torchvision.datasets
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential
from torch.utils.data import DataLoader

dataset = torchvision.datasets.CIFAR10(root="D:\\Python_Project\\pytorch\\dataset2",train=False,transform=torchvision.transforms.ToTensor(),download=True)
dataloader = DataLoader(dataset=dataset,batch_size=1)

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui,self).__init__()
        self.model1 =Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10),
        )

    def forward(self,x):
        x = self.model1(x)
        return x

tudui = Tudui()
# 使用交叉熵损失函数
loss_cross = nn.CrossEntropyLoss()
optim = torch.optim.SGD(tudui.parameters(),lr=0.01)
for epoch in range(20):
    running_loss = 0.0
    for data in dataloader:
        imgs,labels = data
        outputs = tudui(imgs)
        results = loss_cross(outputs,labels)
        optim.zero_grad()
        results.backward()  # 反向传播 这里要注意不能使用定义损失函数那里的 loss，而要使用 调用损失函数之后的 results
        optim.step()
        running_loss = running_loss + results
    print(running_loss)

参考资料：

视频教程：PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】