vgg16模型代码实现pytorch
时间: 2025-01-06 20:44:56 浏览: 8
### VGG16 模型 PyTorch 实现代码示例
为了展示如何使用 PyTorch 实现 VGG16 模型,下面提供了一个完整的代码片段。此代码不仅定义了 VGG16 的架构,还展示了如何加载预训练权重以及如何调整最后几层以适应特定的任务需求。
#### 导入必要的库
```python
import torch
from torchvision import models
import torch.nn as nn
```
#### 定义 VGG16 模型
可以通过 `torchvision.models` 中直接获取预训练的 VGG16 模型:
```python
def get_vgg16_model(num_classes=1000, pretrained=True):
model = models.vgg16(pretrained=pretrained)
# 修改最后一层全连接层以匹配新的分类数目
num_features = model.classifier[6].in_features
features = list(model.classifier.children())[:-1] # 去掉原来的最后一层
features.extend([nn.Linear(num_features, num_classes)]) # 添加新一层
model.classifier = nn.Sequential(*features)
return model
```
这段代码创建了一个函数 `get_vgg16_model()` 来初始化 VGG16 模型,并允许指定输出类别的数量。如果设置了 `pretrained=True` 参数,则会自动下载 ImageNet 上预先训练好的权重[^4]。
#### 加载数据集并验证输入形状
对于给定的数据集,确保输入张量的维度符合预期是非常重要的。根据提供的信息,在训练过程中观察到了如下输入张量的形状:
- 输入图像批次 (X): `[N, C, H, W]`, 即 `[32, 3, 224, 224]`
- 标签向量 (y): 长度为 N (`32`) 的整数标签列表
这些尺寸表明每批包含 32 张 RGB 图像(通道数为 3),分辨率为 224×224 像素。这样的设置非常适合用来作为 VGG16 的输入[^1]。
#### 使用自定义模型进行预测
一旦有了合适的模型和准备就绪的数据集,就可以开始执行前向传播过程来进行推理或评估性能了。这里给出一段简单的测试代码:
```python
if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
# 创建实例化的VGG16模型对象
vgg16 = get_vgg16_model(num_classes=10).to(device)
# 构造模拟输入样本
sample_input = torch.randn(32, 3, 224, 224).to(device)
# 执行一次前向传递得到输出logits
output_logits = vgg16(sample_input)
print(f'Output logits shape: {output_logits.shape}')
```
上述脚本会在 GPU 或 CPU 上运行一个批量大小为 32 的随机输入通过 VGG16 网络,并打印最终输出 logit 向量的形状。注意这里的类别数被设定成了 10 只是为了演示目的;实际应用中应该依据具体任务来决定这个数值[^2]。
阅读全文