def prepare_vgg(model_name, is_pre_trained, fine_tune, num_classes):
if model_name == models.vgg16.__name__:
model = models.vgg16(
pretrained=is_pre_trained,
num_classes=1000 if is_pre_trained else num_classes)
limit_frozen = 5
elif model_name == models.vgg19.__name__:
model = models.vgg19(
pretrained=is_pre_trained,
num_classes=1000 if is_pre_trained else num_classes)
limit_frozen = 7
else:
log.fatal("model name is not known: " + model_name)
sys.exit(1)
if fine_tune:
frozen = nn.Sequential(
*[model.features[i] for i in range(limit_frozen)])
set_parameter_requires_grad(frozen)
model.classifier[-1] = nn.Linear(model.classifier[-1].in_features,
num_classes)
return model
def run_model(model_name, X, y, kf):
collect_garbage()
if model_name == "svm":
run_svm(X=X, y=y, kf=kf)
elif model_name == "lr":
run_lr(X=X, y=y, kf=kf)
elif model_name == "knn":
run_knn(X=X, y=y, kf=kf)
elif model_name == "all":
log.info("Running ML model: svm")
run_svm(X=X, y=y, kf=kf)
log.info("Running ML model: lr")
run_lr(X=X, y=y, kf=kf)
log.info("Running ML model: knn")
run_knn(X=X, y=y, kf=kf)
else:
log.fatal("ML model name is not known: " + model_name)
sys.exit(1)
def get_feature_extractor(model_name, model):
if model_name == models.alexnet.__name__:
feature_extractor = alexnet_feature_extractor(model)
elif model_name in (models.resnet18.__name__, models.resnet50.__name__,
models.resnet152.__name__):
feature_extractor = resnet_feature_extractor(model)
elif model_name in (models.vgg16.__name__, models.vgg19.__name__):
feature_extractor = vgg_feature_extractor(model)
elif model_name == models.densenet169.__name__:
feature_extractor = densenet_feature_extractor(model)
else:
log.fatal("model name is not known: " + model_name)
sys.exit(1)
return feature_extractor
def weighted_model(model_name, pretrain_file, use_actual_num_classes=False):
out_file = ROOT_DIR + "/" + pretrain_file + ".pth"
if model_name == models.alexnet.__name__:
model = models.alexnet(
num_classes=2 if use_actual_num_classes else 1000)
elif model_name == models.resnet18.__name__:
model = models.resnet18(
num_classes=2 if use_actual_num_classes else 1000)
elif model_name == models.resnet50.__name__:
model = models.resnet50(
num_classes=2 if use_actual_num_classes else 1000)
elif model_name == models.resnet152.__name__:
model = models.resnet152(
num_classes=2 if use_actual_num_classes else 1000)
elif model_name == models.vgg16.__name__:
model = models.vgg16(num_classes=2 if use_actual_num_classes else 1000)
elif model_name == models.vgg19.__name__:
model = models.vgg19(num_classes=2 if use_actual_num_classes else 1000)
elif model_name == models.densenet169.__name__:
model = models.densenet169(
num_classes=2 if use_actual_num_classes else 1000)
else:
log.fatal("model name is not known: " + model_name)
sys.exit(1)
try:
log.info("Using class size as: {}".format(
2 if use_actual_num_classes else 1000))
return load_model(model, out_file)
except RuntimeError as re:
log.error(re)
return weighted_model(model_name, pretrain_file, True)
def prepare_resnet(model_name, is_pre_trained, fine_tune, num_classes):
if model_name == models.resnet18.__name__:
model = models.resnet18(
pretrained=is_pre_trained,
num_classes=1000 if is_pre_trained else num_classes)
elif model_name == models.resnet50.__name__:
model = models.resnet50(
pretrained=is_pre_trained,
num_classes=1000 if is_pre_trained else num_classes)
else:
log.fatal("model name is not known: " + model_name)
sys.exit(1)
if fine_tune:
frozen = nn.Sequential(model.conv1, model.bn1, model.relu,
model.maxpool)
set_parameter_requires_grad(frozen)
model.fc = nn.Linear(model.fc.in_features, num_classes)
return model
def main(save=False,
dataset_folder="dataset",
batch_size=20,
img_size=112,
test_without_train=False,
pretrain_file=None,
num_workers=4,
model_name='alexnet',
optimizer_name='Adam',
is_pre_trained=False,
fine_tune=False,
num_epochs=18,
update_lr=True,
normalize=None,
validation_freq=0.1,
lr=0.001,
momentum=0.9,
partial=0.125,
betas=(0.9, 0.99),
weight_decay=0.025):
if not is_pre_trained and fine_tune:
fine_tune = False
if test_without_train and pretrain_file is None:
log.fatal(
"Pretrained weight file is a must on test without train approach")
sys.exit(1)
log.info("Constructing datasets and loaders")
train_data, train_loader, test_data, test_loader = set_dataset_and_loaders(
dataset_folder, batch_size, img_size, num_workers, normalize)
set_0, set_1 = 0, 0
for imgs, labels in test_loader:
if set_0 == 3 and set_1 == 3:
break
for e, label in enumerate(labels.tolist()):
if label == 0 and set_0 != 3:
writer.add_image(
"{} - class image sample {}".format(
train_data.classes[0], set_0),
inv_normalize_tensor(imgs[e], normalize))
set_0 += 1
elif label == 1 and set_1 != 3:
writer.add_image(
"{} - class image sample {}".format(
train_data.classes[1], set_1),
inv_normalize_tensor(imgs[e], normalize))
set_1 += 1
if set_0 == 3 and set_1 == 3:
break
log.info("Calling the model: " + model_name)
if test_without_train:
model = weighted_model(model_name, pretrain_file)
test_model(model, test_loader, 0)
else:
run_model(model_name=model_name,
optimizer_name=optimizer_name,
is_pre_trained=is_pre_trained,
fine_tune=fine_tune,
train_loader=train_loader,
test_loader=test_loader,
num_epochs=num_epochs,
save=save,
update_lr=update_lr,
dataset_folder=dataset_folder,
validation_freq=validation_freq,
lr=lr,
momentum=momentum,
partial=partial,
betas=betas,
weight_decay=weight_decay)
collect_garbage()
writer.close()
def main(transfer_learning, method="", ml_model_name="", cv=10, dataset_folder="dataset",
pretrain_file=None, batch_size=8, img_size=112, num_workers=4, cnn_model_name="", optimizer_name='Adam',
validation_freq=0.1, lr=0.001, momentum=0.9, partial=0.125, betas=(0.9, 0.99), weight_decay=0.025,
update_lr=True, is_pre_trained=False, fine_tune=False, num_epochs=16, normalize=True, seed=17):
if not transfer_learning:
if method.lower() == "ml":
run_ML.main(model_name=ml_model_name, dataset_folder=dataset_folder, seed=seed, cv=cv,
img_size=img_size, normalize=normalize)
elif method.lower() == "cnn":
run_CNN.main(save=False, dataset_folder=dataset_folder, batch_size=batch_size, test_without_train=False,
img_size=img_size, num_workers=num_workers, num_epochs=num_epochs, model_name=cnn_model_name,
optimizer_name=optimizer_name, is_pre_trained=is_pre_trained, fine_tune=fine_tune,
update_lr=update_lr, normalize=normalize, validation_freq=validation_freq, lr=lr,
momentum=momentum, partial=partial, betas=betas, weight_decay=weight_decay)
else:
log.fatal("method name is not known: " + method)
sys.exit(1)
else:
log.info("Constructing datasets and loaders")
train_data, train_loader, test_data, test_loader = set_dataset_and_loaders(dataset_folder=dataset_folder,
batch_size=batch_size,
img_size=img_size,
num_workers=num_workers,
normalize=normalize)
if is_pre_trained and pretrain_file is not None and \
cnn_model_name in pretrain_file.lower():
log.info("Getting PreTrained CNN model: " + cnn_model_name + " from the Weights of " + pretrain_file)
model = cnn_model.weighted_model(cnn_model_name, pretrain_file)
else:
log.info("Running CNN model: " + cnn_model_name)
model = cnn_model.run_model(model_name=cnn_model_name, optimizer_name=optimizer_name, fine_tune=fine_tune,
is_pre_trained=is_pre_trained, train_loader=train_loader, num_epochs=num_epochs,
test_loader=test_loader, validation_freq=validation_freq, lr=lr,
momentum=momentum, partial=partial, betas=betas, weight_decay=weight_decay,
update_lr=update_lr, save=False, dataset_folder=dataset_folder)
log.info("Feature extractor is being created")
feature_extractor = get_feature_extractor(cnn_model_name, model.eval())
log.info("Feature extractor is setting to device: " + str(device))
feature_extractor = feature_extractor.to(device)
log.info("Merging CNN train&test datasets")
dataset = train_data + test_data
log.info("Constructing loader for merged dataset")
data_loader = set_loader(dataset=dataset, batch_size=int(len(dataset) / 5), shuffle=False,
num_workers=num_workers)
log.info("Extracting features as X_cnn array and labels as general y vector")
X_cnn, y = extract_features(data_loader, feature_extractor)
class_dist = {i: y.count(i) for i in y}
class0_size = class_dist[0]
class1_size = class_dist[1]
log.info("Total class 0 size: " + str(class0_size))
log.info("Total class 1 size: " + str(class1_size))
if normalize:
X_cnn = Normalizer().fit_transform(X_cnn)
X_cnn = StandardScaler().fit_transform(X_cnn)
log.info("Number of features in X_cnn: " + str(len(X_cnn[0])))
kf = KFold(n_splits=cv, shuffle=True, random_state=seed)
ml_model.run_model(ml_model_name, X_cnn, y, kf)
collect_garbage()
def run_model(model_name,
optimizer_name,
is_pre_trained,
fine_tune,
train_loader,
test_loader,
validation_freq,
lr,
momentum,
partial,
betas,
weight_decay,
update_lr=True,
num_epochs=25,
save=False,
dataset_folder="dataset"):
collect_garbage()
MODEL_NAME[0] = model_name
num_classes = len(train_loader.dataset.classes)
log.info("Instantiate the model")
if model_name == models.alexnet.__name__:
model = prepare_alexnet(is_pre_trained, fine_tune, num_classes)
elif model_name in (models.resnet18.__name__, models.resnet50.__name__,
models.resnet152.__name__):
model = prepare_resnet(model_name, is_pre_trained, fine_tune,
num_classes)
elif model_name in (models.vgg16.__name__, models.vgg19.__name__):
model = prepare_vgg(model_name, is_pre_trained, fine_tune, num_classes)
elif model_name == models.densenet169.__name__:
model = prepare_densenet(is_pre_trained, fine_tune, num_classes)
else:
log.fatal("model name is not known: " + model_name)
sys.exit(1)
log.info("Setting the model to device")
model = model.to(device)
if "densenet" not in model_name:
log.info("The summary:")
get_summary(model, train_loader)
collect_garbage()
log.info("Setting the loss function")
metric = nn.CrossEntropyLoss()
model_parameters = get_grad_update_params(model, fine_tune)
if optimizer_name == optim.Adam.__name__:
optimizer = optim.Adam(model_parameters, lr=lr)
elif optimizer_name == optim.SGD.__name__:
optimizer = optim.SGD(model_parameters, lr=lr, momentum=momentum)
elif optimizer_name == padam.Padam.__name__:
optimizer = padam.Padam(model_parameters,
lr=lr,
partial=partial,
weight_decay=weight_decay,
betas=betas)
else:
log.fatal("not implemented optimizer name: {}".format(optimizer_name))
sys.exit(1)
log.info("Setting the optimizer as: {}".format(optimizer_name))
SAVE_FILE[0] = (
"" if not is_pre_trained else "PreTrained_"
) + model_name + "_" + optimizer_name + "_" + dataset_folder + "_out.pth"
last_val_iterator = train_model(model,
train_loader,
test_loader,
metric,
optimizer,
lr=lr,
num_epochs=num_epochs,
update_lr=update_lr,
validation_freq=validation_freq,
save=save)
log.info("Testing the model")
test_acc = test_model(model, test_loader, last_val_iterator)
if save and is_verified(test_acc):
exist_files = path_exists(ROOT_DIR, SAVE_FILE[0], "contains")
better = len(exist_files) == 0
if not better:
exist_acc = []
for file in exist_files:
exist_acc.append(float(file.split("_")[0].replace(",", ".")))
better = all(test_acc > acc for acc in exist_acc)
if better:
save_model(model=model,
path=str(round(test_acc, 2)) + "_" + SAVE_FILE[0])
return model