fhipol/deeplearning

1import matplotlib.pyplot as plt2import numpy as np3import os4import torch5import torchvision6import torch.nn as nn7import torch.optim as optim8import torch.optim.lr_scheduler as lr_scheduler9 10from dataset import ObjectDataset11from reporter import ModelReporter12from torch.utils.data import DataLoader13 14from transform import image_transform15 16 17class ModelExecutor:18 19    """20    https://towardsdatascience.com/learning-rate-schedules-and-adaptive-learning-rate-methods-for-deep-learning-2c8f433990d121    https://discuss.pytorch.org/t/where-and-how-to-add-dropout-in-resnet18/1286922    https://stackoverflow.com/questions/47892505/dropout-rate-guidance-for-hidden-layers-in-a-convolution-neural-network23    """24 25    name = ""26 27    # hyperparameters28    train_ratio = 0.729    batch_size = 3230    n_epochs = 1031    learning_rate = 0.0000532    dropout_rate = None33    criterion = nn.CrossEntropyLoss()34 35    model_path = ""36    training_data_path = ""37 38    transform = image_transform39 40    def __init__(self, train_model=True, force_cpu=False) -> None:41 42        self.do_train_model = train_model43        self.model_path = os.getcwd() + "/" + self.model_path44        self.training_data_path = os.getcwd() + "/" + self.training_data_path45 46        # elements47        self.device = torch.device("cuda" if (torch.cuda.is_available() and not force_cpu) else "cpu")48        self.model = self.preload_model()49        self.model.to(self.device)50        print("the device is:", self.device)51 52        self.train_data_loader, self.test_data_loader = self.get_train_and_test_data_loaders() \53                                                            if train_model else (None, None)54 55        # self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate, weight_decay=0.1)56        # self.optimizer = optim.Adam(self.model.parameters(), lr=0.00005, weight_decay=0.1)57        self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate, weight_decay=0.2)58        self.scheduler = lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=self.n_epochs)59 60        # Helpers to store the evolution of the model61        self.reporter = ModelReporter(self.name)62        self.history = {"train": [],63                        "test": []}64 65    def append_dropout(self, model, rate=0.3):66 67        """68        https://discuss.pytorch.org/t/where-and-how-to-add-dropout-in-resnet18/12869/269        the Resnet18 is a pretrained model with a fixed structure70        The quicker way of adding Dropout was patching the model adding the Dropout inside it71        :param rate: the Dropout rate72        :return: the tweaked model73        """74 75        for name, module in model.named_children():76            if len(list(module.children())) > 0:77                self.append_dropout(module)78            if isinstance(module, nn.ReLU):79                print("dropout appended")80                new = nn.Sequential(module, nn.Dropout(p=rate))81                setattr(model, name, new)82        return model83 84    def preload_model(self):85 86        """87        When training, we load the ResNet18 pretrained model and extend it to deal with an output of two classes88        Otherwise, just loading the saved state of the trained (by us) extended model89        :return:90        the model91        """92 93        appended_nn = nn.Sequential(94            nn.Linear(512, 512),95            nn.ReLU(),96            nn.Dropout(p=0.3),97            nn.Linear(512, 256),98            nn.ReLU(),99            nn.Dropout(p=0.3),100            nn.Linear(256, 2),101        )102 103        if self.do_train_model:104            model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)105            # n_last_layer = model.fc.in_features106            model.fc = appended_nn  # nn.Linear(n_last_layer, 2)107            if self.dropout_rate:108                print("setting dropout")109                model = self.append_dropout(model, self.dropout_rate)110            return model111        else:112            print("preloading previous saved model")113            model = torchvision.models.resnet18()114            # n_last_layer = model.fc.in_features115            model.fc = appended_nn  # nn.Linear(n_last_layer, 2)116            model.load_state_dict(torch.load(self.model_path, map_location=self.device))117            return model118 119    def get_train_and_test_data_loaders(self):120 121        dataset = ObjectDataset(self.training_data_path, self.transform)122        train_size = int(len(dataset) * self.train_ratio)123        test_size = len(dataset) - train_size124        train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])125 126        train_data_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)127        test_data_loader = DataLoader(test_dataset, batch_size=self.batch_size, shuffle=True)128        return train_data_loader, test_data_loader129 130    def train_model(self):131 132        print("training model...")133 134        # Loop training:135        for epoch in range(self.n_epochs):136            print(f"\n_____Epoch {epoch}_____ ")137 138            # training phase139            self.execute_epoch_loop(self.train_data_loader, do_training=True)140            self.scheduler.step()141 142            # testing phase143            self.execute_epoch_loop(self.test_data_loader, do_training=False)144 145        # Save the trained model to disk146        torch.save(self.model.state_dict(), self.model_path)147 148    def execute_epoch_loop(self, data_loader, do_training=True):149 150        mode = "train" if do_training else "test"151        loss_records = []152        true_labels = []153        pred_labels = []154 155        if do_training:156            self.model.train()157        else:158            # this should disable dropout layers159            self.model.eval()160 161        for images, labels in data_loader:162 163            # Move the data to the chosen device164            images = images.to(self.device)165            labels = labels.to(self.device)166 167            # output run and loss168            output = self.model(images)169            _, pred = torch.max(output, 1)170 171            loss = self.criterion(output, labels)172            loss_item = loss.item()173            loss_records.append(loss_item)174 175            true_labels.append(labels.cpu().numpy())176            pred_labels.append(pred.cpu().numpy())177 178            if do_training:179 180                self.optimizer.zero_grad()181                loss.backward()182                self.optimizer.step()183 184        print(f"{mode}: In this epoch the loss is: {loss_item}")185 186        this_epoch_history = {f"loss": np.mean(loss_records),187                              f"true_labels": np.concatenate(true_labels),188                              f"pred_labels": np.concatenate(pred_labels), }189 190        # save the data of this iteration on the reporter for reporting at the end of training191        self.reporter.save_data(this_epoch_history, mode)192 193    def run(self):194        if self.do_train_model:195            self.train_model()196            self.reporter.run()197        torch.cuda.empty_cache()198 199 200class HumanModelExecutor(ModelExecutor):201    name = "Human"202    model_path = 'trained_models/trained_model_humans_dev.pth'203    training_data_path = 'humans'204 205 206class BrandsModelExecutor(ModelExecutor):207    name = "Brand"208    model_path = 'trained_models/trained_model_brands_dev.pth'209    training_data_path = 'logos'210 211# if running the script, train both models212if __name__ == "__main__":213 214    model_humans = HumanModelExecutor(train_model=True)215    model_humans.run()216 217    model_logos = BrandsModelExecutor(train_model=True)218    model_logos.run()219 220    plt.show()221