【pytorch】改造model

• Mobilenet_v2

import torchvision.models as models
model=models.mobilenet_v2()

通过查看model的结构，我们需要保留的是model.features，根据需要，修改model.classifier

(classifier): Sequential(
    (0): Dropout(p=0.2, inplace=False)
    (1): Linear(in_features=1280, out_features=1000, bias=True)
  )
 self.model.classifier[0] = nn.Dropout(args.dropout)
 self.model.classifier[1] = nn.Linear(in_features, args.num_class)

• Inception_v3

self.model = models.inception_v3(pretrained=args.pt_imagenet, aux_logits=False)
self.model.fc = nn.Linear(2048, args.num_class)