model.py 文件源码

def __init__(self,embedding_size,num_classes,pretrained=False):
        super(FaceModel, self).__init__()

        self.model = resnet18(pretrained)

        self.embedding_size = embedding_size

        self.model.fc = nn.Linear(512*3*3, self.embedding_size)

        self.model.classifier = nn.Linear(self.embedding_size, num_classes)