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

VGG16是一个经典的卷积神经网络模型，由牛津大学的研究团队开发，特别适用于图像分类任务。在PyTorch中实现VGG16模型的GPU和CPU版本，我们需要加载预训练权重，并区分数据在设备上如何流动。

以下是基本的VGG16模型在PyTorch中使用GPU和CPU的完整代码示例：

# 导入所需库
import torch
import torchvision.models as models
from torchvision import transforms

# GPU或CPU选择
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载预训练的VGG16模型
vgg16 = models.vgg16(pretrained=True)
vgg16.to(device)  # 将模型移到设备上

# 假设我们有一个张量输入
input_tensor = torch.randn(1, 3, 224, 224).to(device)

# 模型前向传播
with torch.no_grad():  # 因为我们不需要梯度计算，所以关闭自动求导
    features = vgg16(input_tensor)

# 如果你想在CPU上运行并查看模型结构，可以这样做：
if device == 'cpu':
    print(vgg16)