pytorch 利用vgg16_bn进行特征提取，然后进行分类

最新推荐文章于 2025-03-05 10:39:16 发布

GeneralJing

最新推荐文章于 2025-03-05 10:39:16 发布

阅读量1.5k

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch

本文链接：https://blog.youkuaiyun.com/GeneralJing/article/details/112602059

pytorch 专栏收录该内容

28 篇文章

订阅专栏

本文介绍了如何使用PyTorch库构建VGG网络模型，包括初始化不同预训练模型、设置CIFAR数据集分类层，以及定义前向传播过程。适合理解深度学习网络结构和迁移学习的应用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

构建网络：

import torch
import torchvision.models as models

class VGG(torch.nn.Module):
    def __init__(self, vgg='vgg16_bn', data_set='CIFAR10', pretrained=False):
        super(VGG, self).__init__()
        self.features = models.__dict__[vgg](pretrained=pretrained).features
        
        classifier = []
        if 'CIFAR' in data_set:
            num_class = int(data_set.split("CIFAR")[1])
            
            classifier.append(torch.nn.Linear(512, 512))
            classifier.append(torch.nn.BatchNorm1d(512))
            classifier.append(torch.nn.Linear(512, num_class))
        else:
            raise RuntimeError("Not expected data flag !!!")

        self.classifier = torch.nn.Sequential(*classifier)
        
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

生成网络：