suyash94/acne_grading

1 2from config import MODEL_DIR, MODEL_INPUT_SIZE, TRANSFORMS_TO_APPLY, MODEL_BACKBONE, MODEL_OBJECTIVE, LAST_N_LAYERS_TO_TRAIN3import os4import torch5import json6from base import TransformationType, ModelBackbone, TrainingObjective7from torchvision import transforms8import torchvision9import torch.nn as nn10 11 12 13def save_model(model, config_json,model_dir = None):14    if model_dir is None:15        model_basedir = MODEL_DIR16        models_present_in_dir = os.listdir(model_basedir)17 18        model_dir_name = 'model_{}'.format(len(models_present_in_dir))19        model_dir = os.path.join(model_basedir, model_dir_name)20        os.mkdir(model_dir)21 22    model_path = os.path.join(model_dir, 'model.pth')23    torch.save(model.state_dict(), model_path)24    config_path = os.path.join(model_dir, 'config.json')25    # import pdb; pdb.set_trace()26    with open(config_path, 'w') as f:27        json.dump(config_json, f)28    29    return model_dir30 31def get_transforms_to_apply_(transformation_type, config_json = None):32    if config_json:33        model_input_size = config_json['MODEL_INPUT_SIZE']34    else:35        model_input_size = MODEL_INPUT_SIZE36 37    if transformation_type == TransformationType.RESIZE:38        return transforms.Resize(model_input_size)39    elif transformation_type == TransformationType.TO_TENSOR:40        return transforms.ToTensor()41    elif transformation_type == TransformationType.RANDOM_HORIZONTAL_FLIP:42        return transforms.RandomHorizontalFlip(p=0.5)43    elif transformation_type == TransformationType.NORMALIZE:44        return transforms.Normalize(mean=[0.485, 0.456, 0.406], 45                                    std=[0.229, 0.224, 0.225])46    elif transformation_type == TransformationType.RANDOM_ROTATION:47        return transforms.RandomRotation(degrees=10)48    elif transformation_type == TransformationType.RANDOM_CLIP:49        return transforms.RandomCrop(model_input_size)50    else:51        raise Exception("Invalid transformation type")52 53def get_transforms_to_apply():54    transforms_to_apply = []55    for transform in TRANSFORMS_TO_APPLY:56        transforms_to_apply.append(get_transforms_to_apply_(TransformationType[transform]))57    return transforms.Compose(transforms_to_apply)58 59def get_model_architecture(config_json = None):60    if config_json:61        model_backbone = ModelBackbone[config_json['MODEL_BACKBONE']]62        model_objective = TrainingObjective[config_json['MODEL_OBJECTIVE']]63    else:64        model_backbone = MODEL_BACKBONE65        model_objective = MODEL_OBJECTIVE66    if model_backbone == ModelBackbone.EFFICIENT_NET_B0:67        if model_objective == TrainingObjective.REGRESSION:68            model = torchvision.models.efficientnet_b0(pretrained=True)69            model.classifier[1] = nn.Sequential(70                nn.Linear(model.classifier[1].in_features, 2048),71                nn.ReLU(),72                nn.Dropout(0.5),73                nn.Linear(2048, 1),74            )75        else:76            raise Exception("Invalid model objective")77    else:78        raise Exception("Invalid model backbone")79    80    return model81 82def get_training_params(model):83    training_params = []84    if MODEL_BACKBONE == ModelBackbone.EFFICIENT_NET_B0:85        if LAST_N_LAYERS_TO_TRAIN > 0:86            for param in model.features[:-LAST_N_LAYERS_TO_TRAIN].parameters():87                param.requires_grad = False88            89            for param in model.features[-LAST_N_LAYERS_TO_TRAIN:].parameters():90                training_params.append(param)91 92 93        for param in model.classifier[1].parameters():94            training_params.append(param)95    else:96        raise Exception("Invalid model backbone")97 98    return training_params99 100def get_criterion():101    if MODEL_OBJECTIVE == TrainingObjective.REGRESSION:102        criterion = nn.MSELoss()103    else:104        raise Exception("Invalid model objective")105    106    return criterion    107      108 109 110 111