我有一个分类问题，需要预测8个类别，例如，我正在使用来自这里的Pytorch中的EfficientNetB3。然而，我对我的自定义类是否正确编写感到困惑。我认为我想去掉预训练模型的最后一层以适应8个输出，对吗？我这样做对了吗？因为当我在我的DataLoader中打印y_preds = model(images)时，它似乎给了我1536个预测。这是预期的行为吗？

!pip install geffnet import geffnetclass EfficientNet(nn.Module):    def __init__(self, config):        super().__init__()        self.config = config        self.model = geffnet.create_model(config.effnet, pretrained=True)        n_features = self.model.classifier.in_features        # fc这个名字重要吗？        self.fc = nn.Linear(n_features, config.num_classes)        self.model.classifier = nn.Identity()            def extract(self, x):        x = self.model(x)        return x    def forward(self, x):        x = self.extract(x).squeeze(-1).squeeze(-1)        return x    model = EfficientNet(config=config)if torch.cuda.is_available():    model.cuda()

打印y_pred的示例代码：

device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)

for step, (images, labels) in enumerate(sample_loader):    images = images.to(device)    labels = labels.to(device)    batch_size = images.shape[0]            y_preds = model(images)    print('4张图像的预测结果如下\n', y_preds)    break

回答：

您甚至在前向传播中没有使用self.fc。

要么只是将其引入为：

def forward(self, x):    ....    x = extract(x)...    x = fc(x)    return x

或者您可以简单地替换名为classifier的层（这样您就不需要Identity层）：

self.model.classifier = nn.Linear(n_features, config.num_classes)

此外，这里config.num_classes应该是8。