【Pytorch框架学习】之分类应用中迁移学习网络修改总结（3）

最新推荐文章于 2023-06-04 07:54:46 发布

一只小金毛zy

最新推荐文章于 2023-06-04 07:54:46 发布

阅读量408

点赞数 1

分类专栏： Pytorch框架

本文链接：https://blog.csdn.net/qq_39071739/article/details/106387359

版权

Pytorch框架专栏收录该内容

6 篇文章 0 订阅

订阅专栏

【Pytorch框架学习】之迁移学习网络分类总结（3）

一、内容

在做图像分类应用时，常常会使用一些经典的预训练网络，比如ResNet、VGG、Inception、DenseNet、EfficientNet、ResNeXt等。但是通常都是在ImageNet上预训练的，所以不能直接使用，因此需要最一些层的修改，这里总结了几种，方法是一样的，这里以二分类为例。
在这里插入图片描述

二、代码

import torch.nn as nn
import torchvision.models as models

num_class = 2

# ResNet
model = models.resnet50()
model.fc = nn.Linear(512, num_class)
print(model)

# AlexNet
model = models.AlexNet()
model.classifier[6] = nn.Linear(4096, num_class)
print(model)

# vgg
model = models.vgg16()
model.classifier[6] = nn.Linear(4096, num_class)
print(model)

# SqueezeNet
model = models.squeezenet1_0()
model.classifier[1] = nn.Conv2d(512, num_class, kernel_size=(1, 1), stride=(1, 1))
print(model)

# DenseNet
model = models.densenet121()
model.classifier = nn.Linear(1024, num_class)
print(model)

# Inception
model = models.inception_v3()
model.AuxLogits.fc = nn.Linear(768, num_class)
model.fc = nn.Linear(2048, num_class)
print(model)