| 40 | |
| 41 | |
| 42 | class VGGBase(nn.Module): |
| 43 | def __init__(self, option='vgg', pret=True, no_pool=False, top=False): |
| 44 | super(VGGBase, self).__init__() |
| 45 | self.dim = 2048 |
| 46 | self.no_pool = no_pool |
| 47 | self.top = top |
| 48 | |
| 49 | if option =='vgg11_bn': |
| 50 | vgg16=models.vgg11_bn(pretrained=pret) |
| 51 | elif option == 'vgg11': |
| 52 | vgg16 = models.vgg11(pretrained=pret) |
| 53 | elif option == 'vgg13': |
| 54 | vgg16 = models.vgg13(pretrained=pret) |
| 55 | elif option == 'vgg13_bn': |
| 56 | vgg16 = models.vgg13_bn(pretrained=pret) |
| 57 | elif option == "vgg16": |
| 58 | vgg16 = models.vgg16(pretrained=pret) |
| 59 | elif option == "vgg16_bn": |
| 60 | vgg16 = models.vgg16_bn(pretrained=pret) |
| 61 | elif option == "vgg19": |
| 62 | vgg16 = models.vgg19(pretrained=pret) |
| 63 | elif option == "vgg19_bn": |
| 64 | vgg16 = models.vgg19_bn(pretrained=pret) |
| 65 | self.classifier = nn.Sequential(*list(vgg16.classifier._modules.values())[:-1]) |
| 66 | self.features = nn.Sequential(*list(vgg16.features._modules.values())[:]) |
| 67 | self.s = nn.Parameter(torch.FloatTensor([10])) |
| 68 | if self.top: |
| 69 | self.vgg = vgg16 |
| 70 | |
| 71 | def forward(self, x, source=True,target=False): |
| 72 | if self.top: |
| 73 | x = self.vgg(x) |
| 74 | return x |
| 75 | else: |
| 76 | x = self.features(x) |
| 77 | x = x.view(x.size(0), 7 * 7 * 512) |
| 78 | x = self.classifier(x) |
| 79 | return x |
| 80 | |
| 81 | |
| 82 | class ResClassifier_MME(nn.Module): |