def setup_model(self, drop_rate):
if train_CNN:
# Initialize modified VGG16
model = models.vgg16(pretrained=True)
model.features[0] = nn.Conv2d(in_channels=1,
out_channels=64,
kernel_size=3,
stride=1,
padding=1)
model.features[4] = nn.Identity()
model.features[16] = nn.Identity()
model.features[23] = nn.Identity()
model.classifier[6] = nn.Linear(in_features=4096,
out_features=10,
bias=True)
model.classifier[2] = nn.Dropout(p=drop_rate)
model.classifier[5] = nn.Dropout(p=drop_rate)
elif train_FC:
# Initialize Fully Connected
model = fullyNet(input_size=28 * 28 * 1,
drop_rate=drop_rate,
init_weights=True)
return model
def setup_model(self, drop_rate):
if train_CNN:
# Loading pretrained ResNext101 model
model = torch.hub.load('facebookresearch/WSL-Images',
'resnext101_32x16d_wsl')
# Setting last fully connected layers to our defined FullyConnected
model.fc = FullyConnected(drop_rate=drop_rate)
elif train_FC:
# Initialize model
model = fullyNet(input_size=224 * 224 * 3,
drop_rate=drop_rate,
num_classes=2)
return model
def setup_model(self):
if train_CNN:
# Initialize modified VGG16
model = models.vgg16(pretrained=True)
model.features[4] = nn.Identity()
model.features[16] = nn.Identity()
model.features[23] = nn.Identity()
model.classifier[6] = nn.Linear(in_features=4096, out_features=10, bias=True)
model.classifier[2] = nn.Dropout(p=0.0)
model.classifier[5] = nn.Dropout(p=0.0)
model.cuda()
elif train_FC:
# Initialize model
model = fullyNet(input_size=32*32*3, drop_rate=0.0, init_weights=True)
model.cuda()
return model