class MyNet(nn.Module): def __init__(self, pretrained=False, **kwargs): super(MyNet, self).__init__() netname = args.arch.split('_')[0] if netname == 'resnet152': self.model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) else: self.model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) if pretrained: parameters = model_zoo.load_url(model_urls[netname]) self.model.load_state_dict(parameters) self.model.avgpool = nn.AvgPool2d(8) self.model.fc = nn.Linear(1024, 1) self.model.sig = nn.Sigmoid() def forward(self, x): x = self.model.conv1(x) x = self.model.bn1(x) x = self.model.relu(x) x = self.model.maxpool(x) x = self.model.layer1(x) x = self.model.layer2(x) x = self.model.layer3(x) # x = self.model.layer4(x) x = self.model.avgpool(x) x = x.view(x.size(0), -1) x = self.model.fc(x) return self.model.sig(x)
def _resnext(arch, num_classes, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: state_dict = load_state_dict_from_url(model_urls[arch], progress=progress) model.load_state_dict(state_dict) in_features = model.fc.in_features model.fc = nn.Linear(in_features=in_features, out_features=num_classes, bias=True) return model
def ResNet9(output_dim=1): model = ResNet(BasicBlock, [1, 1, 1, 1]) in_features = model.fc.in_features model.avgpool = nn.AdaptiveAvgPool2d(1) model.fc = nn.Linear(in_features, output_dim) return model