{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Imports"},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nplt.style.use('bmh')\nimport os\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import models\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import train_test_split","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\ndevice","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Hyperparameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"fine_tune_all_layers = True\nimage_size = 256\nnum_classes = 5\nbatch_size = 32\nnum_epochs = 10\n\nlr = 0.001\nmomentum = 0.9\n    \nfactor = 0.2\npatience = 3","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"model = models.mobilenet_v2(pretrained=True)\n\nif not fine_tune_all_layers:\n    for param in model.parameters():\n        param.requires_grad = False\n        \nin_features = model.classifier[1].in_features\nmodel.classifier[1] = nn.Linear(in_features, num_classes)\nmodel = model.to(device)\n\nmodel_name = 'aug.pt'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Split Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dir = '../input/cassava-leaf-disease-classification/train_images'\ntrain_df = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ids = train_df['image_id'].values\nlabels = train_df['label'].values\n\nX_train_id, X_val_id, y_train, y_val = train_test_split(ids, labels, test_size=0.2, random_state=0, stratify=labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Number of training samples:', len(X_train_id))\nprint('Number of validation samples:', len(X_val_id))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_, counts_train = np.unique(y_train, return_counts=True)\nprint('Training data distribution')\nfor i in range(5):\n  print(\"{}: {:.3f}\".format(i, counts_train[i]/len(y_train)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Dataset Class"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    def __init__(self, data_dir, ids, labels, transform=None):\n        self.data_dir = data_dir\n        self.ids = ids\n        self.labels = labels\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.ids)\n    \n    def __getitem__(self, idx):\n        image = cv2.imread(os.path.join(self.data_dir, self.ids[idx]))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            image = self.transform(image=image)['image']\n        \n        label = self.labels[idx]    \n        \n        return (image, label)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data Augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_transform = A.Compose([\n    A.RandomResizedCrop(256,256),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.25),\n    A.Transpose(p=0.25),\n    A.RandomBrightnessContrast(\n                brightness_limit=(-0.1,0.1), \n                contrast_limit=(-0.1, 0.1), \n                p=0.5),\n    A.Normalize(\n                mean=[0.485, 0.456, 0.406], \n                std=[0.229, 0.224, 0.225], \n                max_pixel_value=255.0, \n                p=1.0),\n    ToTensorV2(p=1.0)\n])\n\n\nvalid_transform = A.Compose([\n    A.Resize(image_size, image_size),\n    A.Normalize(\n                mean=[0.485, 0.456, 0.406], \n                std=[0.229, 0.224, 0.225], \n                max_pixel_value=255.0, \n                p=1.0),\n    ToTensorV2(p=1.0)\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = CassavaDataset(train_dir, X_train_id, y_train, transform=train_transform)\nval_dataset = CassavaDataset(train_dir, X_val_id, y_val, transform=valid_transform)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)\nval_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Loss, Optimizer, and Learning Rate Scheduler"},{"metadata":{"trusted":true},"cell_type":"code","source":"_, counts = np.unique(y_train, return_counts=True)\nbeta = 0.9999\neffective_num = 1.0 - np.power(beta, counts)\nweights = (1 - beta) / effective_num\nweights = weights / np.sum(weights) * num_classes\nweights = weights.astype('float32')\nweights = torch.from_numpy(weights).to(device)\nweights","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# criterion = nn.CrossEntropyLoss(weight=weights)\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=factor, patience=patience, verbose=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train and Validate Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrain_loss_history = [] \nval_loss_history = []\ntrain_acc_history = []\nval_acc_history = []\nbest_val_acc = 0\n\nfor epoch in range(num_epochs):  # loop over the dataset multiple times\n    train_loss = 0.0\n    train_num_correct = 0\n    model.train()\n    for data in train_loader:\n        # get the inputs; data is a list of [inputs, labels]\n        inputs, labels = data\n        inputs = inputs.to(device)\n        labels = labels.to(device)\n\n        # zero the parameter gradients\n        optimizer.zero_grad()\n\n        # forward + backward + optimize\n        outputs = model(inputs)\n        loss = criterion(outputs, labels)\n        train_loss += loss.item()\n        _, pred_class = torch.max(outputs, 1)\n        train_num_correct += pred_class.eq(labels).sum()\n        loss.backward()\n        optimizer.step()\n\n    val_loss = 0.0\n    val_num_correct = 0\n    with torch.no_grad():\n        model.eval()\n        for data in val_loader:\n            inputs, labels = data\n            inputs = inputs.to(device)\n            labels = labels.to(device)\n\n            outputs = model(inputs)\n            loss = criterion(outputs, labels)\n            val_loss += loss.item()\n\n            _, pred_class = torch.max(outputs, 1)\n            val_num_correct += pred_class.eq(labels).sum()\n\n    scheduler.step(val_loss/len(val_loader))\n\n    # get stats\n    train_loss_history.append(train_loss/len(train_loader))\n    val_loss_history.append(val_loss/len(val_loader))\n    train_acc_history.append(train_num_correct/len(train_dataset))\n    val_acc_history.append(val_num_correct/len(val_dataset))\n    if val_acc_history[epoch] > best_val_acc:\n        torch.save(model, model_name)\n        best_val_acc = val_acc_history[epoch]\n    print('Epoch:[{}/{}]\\t Avg Train Loss:{:.3f}\\t Train Acc:{:.3f}\\t Avg Val Loss:{:.3f}\\t Val Acc:{:.3f}\\t'.format(epoch+1, num_epochs, train_loss_history[epoch], \n                                                                              train_acc_history[epoch], val_loss_history[epoch], val_acc_history[epoch]))\n\nprint('Finished Training')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = list(range(1,num_epochs+1))\nplt.plot(epochs, train_loss_history, label='Avg Train Loss')\nplt.plot(epochs, val_loss_history, label='Avg Val loss')\nplt.title(\"Model Loss\")\nplt.ylabel('Loss')\nplt.xlabel('Epoch')\nplt.grid(True)\n# plt.gca().set_ylim(0,1)\nplt.legend();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(epochs, train_acc_history, label='Train Acc')\nplt.plot(epochs, val_acc_history, label='Val Acc')\nplt.title('Model Accuracy')\nplt.ylabel('Accuracy')\nplt.xlabel('Epoch')\nplt.grid(True)\n# plt.gca().set_ylim(0,1)\nplt.legend();\nprint('Best val acc of {:.3f} found at epoch {}'.format(best_val_acc, val_acc_history.index(best_val_acc)+1))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}