def __init__(self, subtype='shufflenet_v2_x1_0', out_stages=[2,3,4], backbone_path=None):
super(ShuffleNetV2, self).__init__()
self.out_stages = out_stages
self.backbone_path = backbone_path
if subtype == 'shufflenet_v2_x0_5':
self.backbone = shufflenet_v2_x0_5(pretrained=not self.backbone_path)
self.out_channels = [24, 48, 96, 192, 1024]
elif subtype == 'shufflenet_v2_x1_0':
self.backbone = shufflenet_v2_x1_0(pretrained=not self.backbone_path)
self.out_channels = [24, 116, 232, 464, 1024]
elif subtype == 'shufflenet_v2_x1_5':
self.backbone = shufflenet_v2_x1_5(pretrained=not self.backbone_path)
self.out_channels = [24, 176, 352, 704, 1024]
elif subtype == 'shufflenet_v2_x2_0':
self.backbone = shufflenet_v2_x2_0(pretrained=not self.backbone_path)
self.out_channels = [24, 244, 488, 976, 2048]
else:
raise NotImplementedError
self.out_channels = self.out_channels[self.out_stages[0]:self.out_stages[-1] + 1]
if self.backbone_path:
self.backbone.load_state_dict(torch.load(self.backbone_path))
else:
self.init_weights()
def get_shufflenet_v2_x1_5(class_num):
model = models.shufflenet_v2_x1_5(pretrained=True)
set_parameter_requires_grad(model)
model.name = 'shufflenet_v2_x1_5'
n_inputs = model.fc.in_features
model.fc = nn.Linear(n_inputs, class_num)
return model, 244
def shufflenet_v2(name, pretrained):
from torchvision.models import shufflenet_v2_x0_5, shufflenet_v2_x1_0, shufflenet_v2_x1_5, shufflenet_v2_x2_0
imagenet_pretrained = pretrained == 'imagenet'
if name == 'shufflenet_v2_x0_5':
net = shufflenet_v2_x0_5(imagenet_pretrained)
elif name == 'shufflenet_v2_x1_0':
net = shufflenet_v2_x1_0(imagenet_pretrained)
elif name == 'shufflenet_v2_x1_5':
net = shufflenet_v2_x1_5(imagenet_pretrained)
elif name == 'shufflenet_v2_x2_0':
net = shufflenet_v2_x2_0(imagenet_pretrained)
else:
raise NotImplementedError(name)
stage1 = nn.Sequential(net.conv1, net.maxpool)
layers = [stage1, net.stage2, net.stage3, net.stage4, net.conv5]
for stage, out in zip(layers, net._stage_out_channels):
stage.out_channels = out
n_pretrained = len(layers) if imagenet_pretrained else 0
return layers, True, n_pretrained
def get_shufflenet_model(
model_name, output_dim,
):
""" Helper function to get a model
:param model_name: Name of the model
:type model_name: str
:param output_dim: Output dimension of the model
:type output_dim: int
:return model: The model
:rtype: ShuffleNetV2
"""
# Getting the model
if model_name == "shufflenet_v2_x0_5":
model = shufflenet_v2_x0_5(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x0_5(pretrained=True)
elif model_name == "shufflenet_v2_x1_0":
model = shufflenet_v2_x1_0(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x1_0(pretrained=True)
elif model_name == "shufflenet_v2_x1_5":
model = shufflenet_v2_x1_5(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x1_5(pretrained=True)
elif model_name == "shufflenet_v2_x2_0":
model = shufflenet_v2_x2_0(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x2_0(pretrained=True)
else:
model = shufflenet_v2_x0_5(num_classes=output_dim)
pretrained_model = models.shufflenet_v2_x0_5(pretrained=True)
IN_FEATURES = pretrained_model.fc.in_features
fc = nn.Linear(IN_FEATURES, output_dim)
pretrained_model.fc = fc
model.load_state_dict(pretrained_model.state_dict())
return model
def __init__(self, model, n_classes, pretrained=True):
super(classifie, self).__init__()
if (model == 'efficientnet-b3'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b3')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b3')
if (model == 'efficientnet-b2'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b2')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b2')
if (model == 'efficientnet-b1'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b1')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b1')
if (model == 'efficientnet-b0'):
if (pretrained == True):
self.cnn_arch = EfficientNet.from_pretrained('efficientnet-b0')
else:
self.cnn_arch = EfficientNet.from_name('efficientnet-b0')
if (model == 'resnet18'):
self.cnn_arch = models.resnet18(pretrained=pretrained)
if (model == 'resnet34'):
self.cnn_arch = models.resnet34(pretrained=pretrained)
if (model == 'resnet50'):
self.cnn_arch = models.resnet50(pretrained=pretrained)
if (model == 'resnet101'):
self.cnn_arch = models.resnet101(pretrained=pretrained)
if (model == 'resnet152'):
self.cnn_arch = models.resnet152(pretrained=pretrained)
if (model == 'densenet121'):
self.cnn_arch = models.densenet121(pretrained=pretrained)
if (model == 'densenet161'):
self.cnn_arch = models.densenet161(pretrained=pretrained)
if (model == 'densenet169'):
self.cnn_arch = models.densenet169(pretrained=pretrained)
if (model == 'densenet201'):
self.cnn_arch = models.densenet201(pretrained=pretrained)
if (model == 'squeezenet1_0'):
self.cnn_arch = models.squeezenet1_0(pretrained=pretrained)
if (model == 'squeezenet1_1'):
self.cnn_arch = models.squeezenet1_1(pretrained=pretrained)
if (model == 'shufflenet_v2_x0_5'):
self.cnn_arch = models.shufflenet_v2_x0_5(pretrained=pretrained)
if (model == 'shufflenet_v2_x1_0'):
self.cnn_arch = models.shufflenet_v2_x1_0(pretrained=pretrained)
if (model == 'shufflenet_v2_x1_5'):
self.cnn_arch = models.shufflenet_v2_x1_5(pretrained=pretrained)
if (model == 'shufflenet_v2_x2_0'):
self.cnn_arch = models.shufflenet_v2_x2_0(pretrained=pretrained)
if (model == 'resnext50_32x4d'):
self.cnn_arch = models.resnext50_32x4d(pretrained=pretrained)
if (model == 'resnext101_32x8d'):
self.cnn_arch = models.resnext101_32x8d(pretrained=pretrained)
self.linear1 = nn.Linear(1000, 256)
self.relu = nn.LeakyReLU()
self.linear2 = nn.Linear(256, n_classes)
#self.softmax = nn.Softmax(dim=1)
self.dropout = nn.Dropout(0.7)
model_ = resnet18(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == 'resnet_34':
model_ = resnet34(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == 'resnet_101':
model_ = resnet101(num_classes=ops.num_classes,
img_size=ops.img_size[0])
elif ops.model == "squeezenet1_0":
model_ = squeezenet1_0(num_classes=ops.num_classes)
elif ops.model == "squeezenet1_1":
model_ = squeezenet1_1(num_classes=ops.num_classes)
elif ops.model == "shufflenetv2":
model_ = ShuffleNetV2(ratio=1., num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x1_5":
model_ = shufflenet_v2_x1_5(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x1_0":
model_ = shufflenet_v2_x1_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x2_0":
model_ = shufflenet_v2_x2_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet":
model_ = ShuffleNet(num_blocks=[2, 4, 2],
num_classes=ops.num_classes,
groups=3)
elif ops.model == "mobilenetv2":
model_ = MobileNetV2(num_classes=ops.num_classes)
elif ops.model == "ReXNetV1":
model_ = ReXNetV1(num_classes=ops.num_classes)
def main(train_loader, train_datasets, valid_loader):
'''args: model_arch, train_loader, train_datasets, valid_loader, epochs, learning_rate, hidden_units, device, save_dir
saves checkpoint under save_dir
returns: checkpoint with {'epoch_tot',
'model',
'criterion',
'optimizer',
'optimizer.state_dict',
'model.state_dict',
'model.class_to_idx'}'''
arg = input_args()
model_arch = arg.arch
hidden_units = arg.hidden_units
learning_rate = arg.learning_rate
device = arg.device
epochs = arg.epochs
save_dir = arg.save_dir
if model_arch == 'alexnet':
gs_vgg = models.alexnet(pretrained=True)
elif model_arch == 'vgg11':
gs_vgg = models.vgg11(pretrained=True)
elif model_arch == 'vgg11_bn':
gs_vgg = models.vgg11_bn(pretrained=True)
elif model_arch == 'vgg13':
gs_vgg = models.vgg13(pretrained=True)
elif model_arch == 'vgg13_bn':
gs_vgg = models.vgg13_bn(pretrained=True)
elif model_arch == 'vgg16':
gs_vgg = models.vgg16(pretrained=True)
elif model_arch == 'vgg16_bn':
gs_vgg = models.vgg16_bn(pretrained=True)
elif model_arch == 'vgg19':
gs_vgg = models.vgg19(pretrained=True)
elif model_arch == 'vgg19_bn':
gs_vgg = models.vgg19_bn(pretrained=True)
elif model_arch == 'resnet18':
gs_vgg = models.resnet18(pretrained=True)
elif model_arch == 'resnet34':
gs_vgg = models.resnet34(pretrained=True)
elif model_arch == 'resnet50':
gs_vgg = models.resnet50(pretrained=True)
elif model_arch == 'resnet101':
gs_vgg = models.resnet101(pretrained=True)
elif model_arch == 'resnet152':
gs_vgg = models.resnet152(pretrained=True)
elif model_arch == 'squeezenet1_0':
gs_vgg = models.squeezenet1_0(pretrained=True)
elif model_arch == 'squeezenet1_1':
gs_vgg = models.squeezenet1_1(pretrained=True)
elif model_arch == 'densenet121':
gs_vgg = models.densenet121(pretrained=True)
elif model_arch == 'densenet169':
gs_vgg = models.densenet169(pretrained=True)
elif model_arch == 'densenet161':
gs_vgg = models.densenet161(pretrained=True)
elif model_arch == 'densenet201':
gs_vgg = models.densenet201(pretrained=True)
elif model_arch == 'inception_v3':
gs_vgg = models.inception_v3(pretrained=True)
elif model_arch == 'googlenet':
gs_vgg = models.googlenet(pretrained=True)
elif model_arch == 'shufflenet_v2_x0_5':
gs_vgg = models.shufflenet_v2_x0_5(pretrained=True)
elif model_arch == 'shufflenet_v2_x1_0':
gs_vgg = models.shufflenet_v2_x1_0(pretrained=True)
elif model_arch == 'shufflenet_v2_x1_5':
gs_vgg = models.shufflenet_v2_x1_5(pretrained=True)
elif model_arch == 'shufflenet_v2_x2_0':
gs_vgg = models.shufflenet_v2_x2_0(pretrained=True)
elif model_arch == 'mobilenet_v2':
gs_vgg = models.mobilenet_v2(pretrained=True)
elif model_arch == 'resnext50_32x4d':
gs_vgg = models.resnext50_32x4d(pretrained=True)
elif model_arch == 'resnext101_32x8d':
gs_vgg = models.resnext101_32x8d(pretrained=True)
elif model_arch == 'wide_resnet50_2':
gs_vgg = models.wide_resnet50_2(pretrained=True)
elif model_arch == 'wide_resnet101_2':
gs_vgg = models.wide_resnet101_2(pretrained=True)
elif model_arch == 'mnasnet0_5':
gs_vgg = models.mnasnet0_5(pretrained=True)
elif model_arch == 'mnasnet0_75':
gs_vgg = models.mnasnet0_75(pretrained=True)
elif model_arch == 'mnasnet1_0':
gs_vgg = models.mnasnet1_0(pretrained=True)
elif model_arch == 'mnasnet1_3':
gs_vgg = models.mnasnet1_3(pretrained=True)
epoch_tot = 0
for parameters in gs_vgg.parameters():
parameters.requires_grad = False
try:
input_layer = gs_vgg.classifier[0].in_features
except:
input_layer = gs_vgg.classifier.in_features
hidden_layers = [(int(hidden_units * 0.68)), (int(hidden_units * 0.32))]
output_layer = len(train_loader)
gs_vgg.classifier = nn.Sequential(
nn.Linear(input_layer, hidden_layers[0]),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(hidden_layers[0], hidden_layers[1]),
nn.ReLU(),
#nn.Linear(hidden_layers[1], output_layer),
nn.Linear(hidden_layers[1], 102),
nn.LogSoftmax(dim=1))
criterion = nn.NLLLoss()
optimizer = optim.Adam(gs_vgg.classifier.parameters(), lr=learning_rate)
gs_vgg.to(device)
step_num = 0
epoch = epochs
running_loss = 0
print_every = 10
for e in range(epoch):
epoch_tot += 1
for images, labels in train_loader:
gs_vgg.train()
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
output = gs_vgg.forward(images)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
step_num += 1
if step_num % print_every == 0:
test_loss = 0
accuracy = 0
gs_vgg.eval()
for images, labels in valid_loader:
images, labels = images.to(device), labels.to(device)
output = gs_vgg.forward(images)
loss = criterion(output, labels)
test_loss += loss.item()
prob = torch.exp(output)
top_p, top_class = prob.topk(1, dim=1)
equals = top_class == labels.view(*top_class.shape)
accuracy += torch.mean(equals.type(
torch.FloatTensor)).item()
print(
f"Total Epochs: {epoch_tot}.. "
f"Train loss: {running_loss/print_every:.3f}.. "
f"Test loss: {test_loss/len(valid_loader):.3f}.. "
f"Test accuracy: {(accuracy/len(valid_loader))*100:.1f}%")
running_loss = 0
gs_vgg.train()
gs_vgg.class_to_idx = train_datasets.class_to_idx
gs_checkpoint = {
'epoch_tot': epoch_tot,
'model': gs_vgg,
'criterion': criterion,
'optimizer': optimizer,
'optimizer.state_dict': optimizer.state_dict(),
'model.state_dict': gs_vgg.state_dict(),
'model.class_to_idx': gs_vgg.class_to_idx
}
torch.save(gs_checkpoint, save_dir)
return gs_checkpoint
def get_model(name, device):
if name == "normal_cnn":
model = Net().to(device)
return model
elif name == "alexnet":
model = models.alexnet().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg11":
model = models.vgg11().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg13":
model = models.vgg13().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg16":
model = models.vgg16().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg19":
model = models.vgg19().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg11_bn":
model = models.vgg11_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg13_bn":
model = models.vgg13_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg16_bn":
model = models.vgg16_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "vgg19_bn":
model = models.vgg19_bn().to(device)
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet121":
model = models.densenet121().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet161":
model = models.densenet161().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet169":
model = models.densenet169().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "densenet201":
model = models.densenet201().to(device)
num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet18":
model = models.resnet18().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet34":
model = models.resnet34().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet50":
model = models.resnet50().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet101":
model = models.resnet101().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnet152":
model = models.resnet152().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "googlenet":
model = models.googlenet(aux_logits=False).to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "inception_v3":
model = models.inception_v3(aux_logits=False).to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x0_5":
model = models.shufflenet_v2_x0_5().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x1_0":
model = models.shufflenet_v2_x1_0().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x1_5":
model = models.shufflenet_v2_x1_5().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "shufflenet_v2_x2_0":
model = models.shufflenet_v2_x2_0().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "mobilenet_v2":
model = models.mobilenet_v2().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnext50_32x4d":
model = models.resnext50_32x4d().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "resnext101_32x8d":
model = models.resnext101_32x8d().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "wide_resnet50_2":
model = models.wide_resnet50_2().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "wide_resnet101_2":
model = models.wide_resnet101_2().to(device)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet0_5":
model = models.mnasnet0_5().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet0_75":
model = models.mnasnet0_75().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet1_0":
model = models.mnasnet1_0().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
elif name == "mnasnet1_3":
model = models.mnasnet1_3().to(device)
num_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_features, 1).to(device)
return model
def get_model(model_id, use_pretrained):
model_ft = None
if model_id == PyTorchModelsEnum.ALEXNET:
model_ft = models.alexnet(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.DENSENET121:
model_ft = models.densenet121(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.DENSENET161:
model_ft = models.densenet161(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.DENSENET169:
model_ft = models.densenet169(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.DENSENET201:
model_ft = models.densenet201(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.GOOGLENET:
model_ft = models.googlenet(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.INCEPTION_V3:
model_ft = models.inception_v3(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.MOBILENET_V2:
model_ft = models.mobilenet_v2(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.MNASNET_0_5:
model_ft = models.mnasnet0_5(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.MNASNET_0_75: # no pretrained
model_ft = models.mnasnet0_75(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.MNASNET_1_0:
model_ft = models.mnasnet1_0(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.MNASNET_1_3:
model_ft = models.mnasnet1_3(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNET18:
model_ft = models.resnet18(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNET34:
model_ft = models.resnet34(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNET50:
model_ft = models.resnet50(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNET101:
model_ft = models.resnet101(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNET152:
model_ft = models.resnet152(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNEXT50:
model_ft = models.resnext50_32x4d(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.RESNEXT101:
model_ft = models.resnext101_32x8d(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.SHUFFLENET_V2_0_5:
model_ft = models.shufflenet_v2_x0_5(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.SHUFFLENET_V2_1_0:
model_ft = models.shufflenet_v2_x1_0(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.SHUFFLENET_V2_1_5:
model_ft = models.shufflenet_v2_x1_5(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.SHUFFLENET_V2_2_0:
model_ft = models.shufflenet_v2_x2_0(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.SQUEEZENET1_0:
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.SQUEEZENET1_1:
model_ft = models.squeezenet1_1(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG11:
model_ft = models.vgg11(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG11_BN:
model_ft = models.vgg11_bn(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG13:
model_ft = models.vgg13(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG13_BN:
model_ft = models.vgg13_bn(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG16:
model_ft = models.vgg16(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG16_BN:
model_ft = models.vgg16_bn(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG19:
model_ft = models.vgg19(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.VGG19_BN:
model_ft = models.vgg19_bn(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.WIDE_RESNET50:
model_ft = models.wide_resnet50_2(pretrained=use_pretrained)
elif model_id == PyTorchModelsEnum.WIDE_RESNET101:
model_ft = models.wide_resnet101_2(pretrained=use_pretrained)
return model_ft
def initialModel(modelName, numClasses, featureExtract, usePretrained):
if modelName == "vgg19":
modelUse = models.vgg19(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0'] = nn.Conv2d(1,
64,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1))
modelUse.classifier._modules['6'] = nn.Linear(in_features=4096,
out_features=numClasses,
bias=True)
elif modelName == "shufflenet_v2_x2_0":
modelUse = models.shufflenet_v2_x2_0(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(2048, numClasses, bias=True))
else:
modelUse.fc = nn.Linear(2048, numClasses, bias=True)
elif modelName == "shufflenet_v2_x0_5":
modelUse = models.shufflenet_v2_x0_5(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(
nn.Dropout(dropNum),
nn.Linear(1024, numClasses, bias=True),
)
else:
modelUse.fc = nn.Linear(1024, numClasses, bias=True)
elif modelName == "shufflenet_v2_x1_5":
modelUse = models.shufflenet_v2_x1_5(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(
nn.Dropout(dropNum),
nn.Linear(1024, numClasses, bias=True),
)
else:
modelUse.fc = nn.Linear(1024, numClasses, bias=True)
elif modelName == "shufflenet_v2_x1_0":
modelUse = models.shufflenet_v2_x1_0(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(1024, numClasses, bias=True))
else:
modelUse.fc = nn.Linear(1024, numClasses, bias=True)
elif modelName == "squeezenet1_1":
modelUse = models.squeezenet1_1(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0'] = nn.Conv2d(1,
64,
kernel_size=(3, 3),
stride=(2, 2))
modelUse.classifier = nn.Sequential(
nn.Dropout(p=dropNum, inplace=False),
nn.Conv2d(512, numClasses, kernel_size=(1, 1), stride=(1, 1)),
nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(output_size=(1, 1)))
elif modelName == "densenet121":
modelUse = models.densenet121(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['conv0'] = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.classifier = nn.Sequential(
nn.Dropout(dropNum), nn.Linear(1024, numClasses, bias=True))
else:
modelUse.classifier = nn.Linear(1024, numClasses, bias=True)
elif modelName == "vgg16":
modelUse = models.vgg16(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0'] = nn.Conv2d(1,
64,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1))
modelUse.classifier._modules['6'] = nn.Linear(in_features=4096,
out_features=numClasses,
bias=True)
elif modelName == 'mobilenet_v2':
modelUse = models.mobilenet_v2(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0']._modules['0'] = nn.Conv2d(
1,
32,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
modelUse.classifier = nn.Sequential(
nn.Dropout(p=dropNum, inplace=False),
nn.Linear(in_features=1280, out_features=numClasses, bias=True))
elif modelName == "resnet18":
modelUse = models.resnet18(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
numFeatures = modelUse.fc.in_features
if mono:
modelUse.conv1 = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(numFeatures, numClasses))
else:
modelUse.fc = nn.Linear(numFeatures, numClasses)
elif modelName == "resnet34":
modelUse = models.resnet34(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
numFeatures = modelUse.fc.in_features
if mono:
modelUse.conv1 = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(numFeatures, numClasses))
else:
modelUse.fc = nn.Linear(numFeatures, numClasses)
elif modelName == "resnet50":
modelUse = models.resnet50(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
numFeatures = modelUse.fc.in_features
if mono:
modelUse.conv1 = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(numFeatures, numClasses))
else:
modelUse.fc = nn.Linear(numFeatures, numClasses)
else:
print("model not implemented")
return None, None
return modelUse
def trainer(ops, f_log):
try:
os.environ['CUDA_VISIBLE_DEVICES'] = ops.GPUS
if ops.log_flag:
sys.stdout = f_log
set_seed(ops.seed)
#---------------------------------------------------------------- 构建模型
if ops.model == 'resnet_50':
model_ = resnet50(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == 'resnet_18':
model_ = resnet18(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == 'resnet_34':
model_ = resnet34(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == 'resnet_101':
model_ = resnet101(pretrained=True,
num_classes=ops.num_classes,
img_size=ops.img_size[0],
dropout_factor=ops.dropout)
elif ops.model == "squeezenet1_0":
model_ = squeezenet1_0(pretrained=True,
num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "squeezenet1_1":
model_ = squeezenet1_1(pretrained=True,
num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "shufflenetv2":
model_ = ShuffleNetV2(ratio=1.,
num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "shufflenet_v2_x1_5":
model_ = shufflenet_v2_x1_5(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x1_0":
model_ = shufflenet_v2_x1_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet_v2_x2_0":
model_ = shufflenet_v2_x2_0(pretrained=False,
num_classes=ops.num_classes)
elif ops.model == "shufflenet":
model_ = ShuffleNet(num_blocks=[2, 4, 2],
num_classes=ops.num_classes,
groups=3,
dropout_factor=ops.dropout)
elif ops.model == "mobilenetv2":
model_ = MobileNetV2(num_classes=ops.num_classes,
dropout_factor=ops.dropout)
elif ops.model == "ReXNetV1":
model_ = ReXNetV1(num_classes=ops.num_classes,
dropout_factor=ops.dropout)
else:
print(" no support the model")
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
model_ = model_.to(device)
# print(model_)# 打印模型结构
# Dataset
dataset = LoadImagesAndLabels(ops=ops,
img_size=ops.img_size,
flag_agu=ops.flag_agu,
fix_res=ops.fix_res,
vis=False)
print("handpose done")
print('len train datasets : %s' % (dataset.__len__()))
# Dataloader
dataloader = DataLoader(dataset,
batch_size=ops.batch_size,
num_workers=ops.num_workers,
shuffle=True,
pin_memory=False,
drop_last=True)
# 优化器设计
optimizer_Adam = torch.optim.Adam(model_.parameters(),
lr=ops.init_lr,
betas=(0.9, 0.99),
weight_decay=1e-6)
# optimizer_SGD = optim.SGD(model_.parameters(), lr=ops.init_lr, momentum=ops.momentum, weight_decay=ops.weight_decay)# 优化器初始化
optimizer = optimizer_Adam
# 加载 finetune 模型
if os.access(ops.fintune_model, os.F_OK): # checkpoint
chkpt = torch.load(ops.fintune_model, map_location=device)
model_.load_state_dict(chkpt)
print('load fintune model : {}'.format(ops.fintune_model))
print('/**********************************************/')
# 损失函数
if ops.loss_define != 'wing_loss':
criterion = nn.MSELoss(reduce=True, reduction='mean')
step = 0
idx = 0
# 变量初始化
best_loss = np.inf
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
flag_change_lr_cnt = 0 # 学习率更新计数器
init_lr = ops.init_lr # 学习率
epochs_loss_dict = {}
for epoch in range(0, ops.epochs):
if ops.log_flag:
sys.stdout = f_log
print('\nepoch %d ------>>>' % epoch)
model_.train()
# 学习率更新策略
if loss_mean != 0.:
if best_loss > (loss_mean / loss_idx):
flag_change_lr_cnt = 0
best_loss = (loss_mean / loss_idx)
else:
flag_change_lr_cnt += 1
if flag_change_lr_cnt > 50:
init_lr = init_lr * ops.lr_decay
set_learning_rate(optimizer, init_lr)
flag_change_lr_cnt = 0
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
for i, (imgs_, pts_) in enumerate(dataloader):
# print('imgs_, pts_',imgs_.size(), pts_.size())
if use_cuda:
imgs_ = imgs_.cuda(
) # pytorch 的 数据输入格式 ： (batch, channel, height, width)
pts_ = pts_.cuda()
output = model_(imgs_.float())
if ops.loss_define == 'wing_loss':
loss = got_total_wing_loss(output, pts_.float())
else:
loss = criterion(output, pts_.float())
loss_mean += loss.item()
loss_idx += 1.
if i % 10 == 0:
loc_time = time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime())
print(' %s - %s - epoch [%s/%s] (%s/%s):'%(loc_time,ops.model,epoch,ops.epochs,i,int(dataset.__len__()/ops.batch_size)),\
'Mean Loss : %.6f - Loss: %.6f'%(loss_mean/loss_idx,loss.item()),\
' lr : %.8f'%init_lr,' bs :',ops.batch_size,\
' img_size: %s x %s'%(ops.img_size[0],ops.img_size[1]),' best_loss: %.6f'%best_loss)
# 计算梯度
loss.backward()
# 优化器对模型参数更新
optimizer.step()
# 优化器梯度清零
optimizer.zero_grad()
step += 1
torch.save(
model_.state_dict(),
ops.model_exp + '{}-size-{}-model_epoch-{}.pth'.format(
ops.model, ops.img_size[0], epoch))
except Exception as e:
print('Exception : ', e) # 打印异常
print('Exception file : ',
e.__traceback__.tb_frame.f_globals['__file__']) # 发生异常所在的文件
print('Exception line : ', e.__traceback__.tb_lineno) # 发生异常所在的行数
def Model(self):
if self.arch == 'alexnet':
model = models.alexnet(num_classes=self.num_classes)
if self.arch == 'vgg11':
model = models.vgg11(num_classes=self.num_classes)
if self.arch == 'vgg13':
model = models.vgg13(num_classes=self.num_classes)
if self.arch == 'vgg16':
model = models.vgg16(num_classes=self.num_classes)
if self.arch == 'vgg19':
model = models.vgg19(num_classes=self.num_classes)
if self.arch == 'vgg11_bn':
model = models.vgg11_bn(num_classes=self.num_classes)
if self.arch == 'vgg13_bn':
model = models.vgg13_bn(num_classes=self.num_classes)
if self.arch == 'vgg16_bn':
model = models.vgg16_bn(num_classes=self.num_classes)
if self.arch == 'vgg19_bn':
model = models.vgg19_bn(num_classes=self.num_classes)
if self.arch == 'resnet18':
model = models.resnet18(num_classes=self.num_classes)
if self.arch == 'resnet34':
model = models.resnet34(num_classes=self.num_classes)
if self.arch == 'resnet50':
model = models.resnet50(num_classes=self.num_classes)
if self.arch == 'resnet101':
model = models.resnet101(num_classes=self.num_classes)
if self.arch == 'resnet152':
model = models.resnet152(num_classes=self.num_classes)
if self.arch == 'squeezenet1_0':
model = models.squeezenet1_0(num_classes=self.num_classes)
if self.arch == 'squeezenet1_1':
model = models.squeezenet1_1(num_classes=self.num_classes)
if self.arch == 'densenet121':
model = models.densenet121(num_classes=self.num_classes)
if self.arch == 'densenet161':
model = models.densenet161(num_classes=self.num_classes)
if self.arch == 'densenet169':
model = models.densenet169(num_classes=self.num_classes)
if self.arch == 'densenet201':
model = models.densenet201(num_classes=self.num_classes)
if self.arch == 'inception_v1':
# parameters 'aux_logits' maybe will make the model not work
model = models.googlenet(num_classes=self.num_classes)
if self.arch == 'inception_v3':
# parameters 'aux_logits' maybe will make the model not work
model = models.inception_v3(num_classes=self.num_classes)
if self.arch == 'shufflenet_v2_x0_5':
model = models.shufflenet_v2_x0_5(num_classes=self.num_classes)
if self.arch == 'shufflenet_v2_x1_0':
model = models.shufflenet_v2_x1_0(num_classes=self.num_classes)
if self.arch == 'shufflenet_v2_x1_5':
model = models.shufflenet_v2_x1_5(num_classes=self.num_classes)
if self.arch == 'shufflenet_v2_x2_0':
model = models.shufflenet_v2_x2_0(num_classes=self.num_classes)
if self.arch == 'mobilenet_v2':
model = models.mobilenet_v2(num_classes=self.num_classes)
if self.arch == 'resnext50_32x4d':
model = models.resnext50_32x4d(num_classes=self.num_classes)
if self.arch == 'resnext101_32x4d':
model = models.resnext101_32x4d(num_classes=self.num_classes)
if self.arch == 'wide_resnet50_2':
model = models.wide_resnet50_2(num_classes=self.num_classes)
if self.arch == 'wide_resnet101_2':
model = models.wide_resnet101_2(num_classes=self.num_classes)
if self.arch == 'mnasnet1_0':
model = models.mnasnet1_0(num_classes=self.num_classes)
model = torch.nn.DataParallel(model, device_ids=self.gups).cuda()
return model
vgg13_bn = models.vgg13_bn(pretrained=True) vgg16_bn = models.vgg16_bn(pretrained=True) vgg19 = models.vgg19(pretrained=True) vgg19_bn = models.vgg19(pretrained=True) resnet34 = models.resnet34(pretrained=True) resnet50 = models.resnet50(pretrained=True) resnet101 = models.resnet101(pretrained=True) resnet152 = models.resnet152(pretrained=True) squeezenet1_0 = models.squeezenet1_0(pretrained=True) squeezenet1_1 = models.squeezenet1_1(pretrained=True) densenet121 = models.densenet121(pretrained=True) densenet161 = models.densenet161(pretrained=True) densenet169 = models.densenet169(pretrained=True) densenet201 = models.densenet201(pretrained=True) inception_v3 = models.inception_v3(pretrained=True) shufflenet_v2_x0_5 = models.shufflenet_v2_x0_5(pretrained=True) shufflenet_v2_x1_0 = models.shufflenet_v2_x1_0(pretrained=True) shufflenet_v2_x1_5 = models.shufflenet_v2_x1_5(pretrained=True) shufflenet_v2_x2_0 = models.shufflenet_v2_x2_0(pretrained=True) mobilenet_v2 = models.mobilenet_v2(pretrained=True) resnext50_32x4d = models.resnext50_32x4d(pretrained=True) resnext101_32x8d = modelos.resnext101_32x8d(pretrained=True) # vgg = models.vgg16(pretrained = True) # googlenet = models.googlenet(pretrained = True) # resnet18 = models.resnet18(pretrained = True) # model_info = [(alexnet, 'AlexNet'), (vgg, 'VGG16'), (googlenet, 'GoogLeNet'), (resnet18, 'ResNet18')] # plot_num_weights(model_info)