基于pytorch的vgg16代码

以下是基于PyTorch的VGG16代码：

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

class VGG16(nn.Module):
    def __init__(self, num_classes):
        super(VGG16, self).__init__()
        self.features = models.vgg16(pretrained=True).features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )

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

该代码使用了预训练的VGG16模型，并在其基础上添加了一个全连接层作为分类器。其中，`num_classes`即为分类数，可以根据具体的任务进行调整。在`forward`函数中，输入数据首先通过VGG16的卷积层和池化层进行特征提取，然后通过全连接层进行分类。