VGG16模型 pytorch GPU和CPU版 完整代码

VGG16是一种经典的卷积神经网络模型，由Karen Simonyan和Andrew Zisserman于2014年在ImageNet竞赛中提出。它以其深度（16个卷积层）而闻名，且结构简单，易于理解和实现。

在PyTorch中，使用GPU版本的VGG16通常会更快，因为GPU并行计算能力强，适合大规模数据的处理。下面是使用预训练的VGG16模型（包括GPU和CPU版本）的基本代码示例：

# 导入必要的库
import torch
from torchvision import models, transforms

# 使用GPU（如果可用）
if torch.cuda.is_available():
    device = torch.device('cuda')  # 默认选择GPU
else:
    device = torch.device('cpu')

# 加载预训练的VGG16模型（默认去掉最后的分类层）
model_vgg16 = models.vgg16(pretrained=True)
model_vgg16.to(device)

# 如果你想在CPU上运行，可以将模型转移到CPU
# model_vgg16.to(torch.device('cpu'))

# 示例使用：
# 图片预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载图片
input_image = transform(image).unsqueeze(0)  # unsqueeze是为了匹配模型期望的输入形状

# 将图像移动到设备
input_image = input_image.to(device)

# 运行前向传播
with torch.no_grad():  # 防止梯度跟踪
    output = model_vgg16(input_image)