关于ImageNet预训练分类模型的构建

[itisianlee]

class CLSHead(nn.Module):
    def __init__(self, in_ch=384, out_ch=1000, dropout=0.2):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.bn = nn.BatchNorm2d(in_ch)
        self.dropout = nn.Dropout(p=dropout, inplace=True)
        self.l = nn.Linear(in_ch, out_ch)

    def forward(self, x):
        x = self.bn(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.l(x)
        return x

分类的头是这样的写的吗？是直接接到：self.trans = BasicLayer的输出上吗？ 我使用google drive上checkpoint加载评测ImageNet val精度是59.5左右。 期待回复，谢谢