{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"# original author: github.com/schatt89/","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Helper files**"},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nutils.py\n\"\"\"\n\nimport matplotlib.pyplot as plt\nimport os\nfrom shutil import copytree, ignore_patterns\nimport numpy as np\n\nfrom typing import Dict, Tuple\n\n\ndef workdir_copy(pwd: str, copy_path: str):\n    cp_path = os.path.join(copy_path, 'wdir_copy')\n    copytree(pwd, cp_path, ignore=ignore_patterns('__pycache__', '.git'))\n\ndef save_predictions(\n        save_path: str,\n        predictions: Tuple[str, int],\n        idx_to_class: Dict[int, str]\n    ) -> None:\n    '''\n        Format:\n        Id,Category\n        0,Car\n        1,Catepillar\n    '''\n    with open(save_path, 'w') as outf:\n        # header\n        outf.write('Id,Category\\n')\n\n        # other lines\n        for (pred_path, pred_idx) in predictions:\n            # extract Id from the filename\n            Id = int(os.path.split(pred_path)[1].strip('.jpg'))\n            outf.write(f'{Id},{idx_to_class[pred_idx]}\\n')\n\n    print(f'Wrote preds to {save_path}')\n\n\ndef plot_images(images, data_dir, cls_true, cls_pred=None):\n    \"\"\"\n    Adapted from https://github.com/Hvass-Labs/TensorFlow-Tutorials/\n    \"\"\"\n    label_names = sorted(os.listdir(data_dir))\n    fig, axes = plt.subplots(5, 8, figsize=(15, 10))\n\n    for i, ax in enumerate(axes.flat):\n        # plot img\n        means = np.array([0.485, 0.456, 0.406])\n        stds = np.array([0.229, 0.224, 0.225])\n\n        inp = np.clip(images[i, :, :, :] * stds + means, 0, 1)\n        ax.imshow(inp, interpolation='spline16')\n\n        # show true & predicted classes\n        cls_true_name = label_names[cls_true[i]]\n        if cls_pred is None:\n            xlabel = \"{0} ({1})\".format(cls_true_name, cls_true[i])\n        else:\n            cls_pred_name = label_names[cls_pred[i]]\n            xlabel = \"True: {0}\\nPred: {1}\".format(\n                cls_true_name, cls_pred_name\n            )\n        ax.set_xlabel(xlabel)\n        ax.set_xticks([])\n        ax.set_yticks([])\n\n    plt.show(block=False)\n    plt.pause(10)\n    plt.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nsamplers.py\n\"\"\"\n\nimport numpy as np\nimport torch\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import Dataset\nfrom torch.utils.data.sampler import Sampler, SubsetRandomSampler, WeightedRandomSampler\nfrom typing import Dict, Tuple, Iterator, Iterable\n\n\ndef sklearn_stratified_train_valid_splits(\n    targets: np.ndarray,\n    valid_size: float,\n    random_state: int\n) -> Iterator[Tuple[list, list]]:\n\n    # X which is not used in .split anyway\n    X = np.arange(len(targets)).reshape(-1, 1)\n    # n_splits=int(1/0.25): 4\n    skf = StratifiedKFold(n_splits=int(1/valid_size), shuffle=True, random_state=random_state)\n    return skf.split(X, targets)\n\n\ndef get_train_class_weights(targets: np.ndarray, train_index: Iterable) -> Dict[int, float]:\n    # filter targets only those that belong to training set\n    train_targets = [targets[idx] for idx in train_index]\n    # count each class only in training; if valid is used -> may overfit\n    cls_to_count = {cls: train_targets.count(cls) for cls in set(train_targets)}\n    # weights are disproportional to the number of occurences for each class in train dataset\n    cls_to_weight = {cls: 1 / cls_to_count[cls] for cls, count in cls_to_count.items()}\n    return cls_to_weight\n\n\ndef calculate_samples_weights(\n    train_index: list,\n    targets: np.ndarray\n) -> Tuple[torch.DoubleTensor, Dict[int, float]]:\n\n    # dict: idx -> weight\n    cls_to_weight = get_train_class_weights(targets, train_index)\n    # initialize the weights with zeros\n    samples_weights = torch.zeros(len(targets)).double()\n    # fill all training indices with weights according to the sample size; valid idx have probs = 0\n    for idx in train_index:\n        sample_class = targets[idx]\n        samples_weights[idx] = cls_to_weight[sample_class]\n\n    return samples_weights, cls_to_weight\n\ndef valid_and_train_samplers(\n    train_dataset: Dataset,\n    weighted: bool,\n    valid_size: float,\n    random_state: int,\n) -> Tuple[Sampler, Sampler, Dict[int, float]]:\n\n    # extract only the classes and make them numpy array\n    targets = np.array(train_dataset.targets)\n    # train-valid split in a stratified manner (return splits iterator)\n    skf_splits = sklearn_stratified_train_valid_splits(targets, valid_size, random_state)\n    # select only the first split (can be expanded to K-fold cross validation)\n    train_index, valid_index = next(skf_splits)\n\n    if weighted:\n        print('Using weighted sampler')\n        # weights for each sample in ImageDataset (train: weighted, valid: all zeros); + class -> weight\n        samples_weights, cls_to_weight = calculate_samples_weights(train_index, targets)\n        # sampled indices (which must belong only to train_index) <-\n        train_sampler = WeightedRandomSampler(samples_weights, len(train_index))\n    else:\n        # dict: idx -> weight\n        cls_to_weight = get_train_class_weights(targets, train_index)\n        # non-weighted random sampler\n        train_sampler = SubsetRandomSampler(train_index)\n\n    # cls_to_count_obs = {cls: 0 for cls in set(targets)}\n    # for i in range(100):\n    #     for train_idx in train_sampler:\n    #         cls_to_count_obs[targets[train_idx]] += 1\n    # print(cls_to_count_obs) # {0: 2501153, 1: 2497812, 2: 2501595, 3: 2499540}\n    \n    return train_sampler, SubsetRandomSampler(valid_index), cls_to_weight","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Dataset**"},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\ntransforms.py \n\"\"\"\n\nimport numpy as np\nfrom imgaug import augmenters as iaa\n\nclass ImgAugTransform(object):\n    \n    def __init__(self, model_input, p=0.25):\n        self.model_input = model_input\n\n        self.aug = iaa.Sequential([\n            iaa.Sometimes(p, \n                iaa.OneOf([\n                    iaa.MotionBlur(k=15, angle=[-135, -90, -45, 45, 90, 135]),\n                    iaa.GaussianBlur(sigma=(0, 3.0)),\n                ])\n            ),\n            iaa.Sometimes(p, iaa.Affine(rotate=(-20, 20), mode='symmetric')),\n            iaa.Sometimes(p, iaa.AddToHueAndSaturation(value=(-10, 10), per_channel=True)),\n            iaa.Sometimes(p, iaa.ContrastNormalization((0.5, 1.5), per_channel=0.5)),\n            iaa.Sometimes(p, iaa.Sharpen(alpha=0.5)),\n            iaa.Sometimes(p, iaa.AdditiveGaussianNoise(scale=(0, 0.1*255))),\n            iaa.Sometimes(p, iaa.Add((-20, 20), per_channel=0.5)),\n            iaa.Sometimes(p, iaa.PiecewiseAffine(scale=(0.01, 0.02))),\n            iaa.Sometimes(p, iaa.PerspectiveTransform(scale=(0.01, 0.15))),\n            # iaa.ChannelShuffle(p/5)\n            \n        ])\n\n    def __call__(self, img):\n        img = np.array(img)\n        return self.aug.augment_image(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\ndataset.py\n\"\"\"\n\nimport os\nfrom glob import glob\n\nfrom PIL import Image\nfrom torch.utils.data import Dataset\n\n\nclass TestImageFolder(Dataset):\n\n    def __init__(self, data_dir, transforms):\n        self.data_dir = data_dir\n        self.dataset = sorted(glob(os.path.join(data_dir, '*.jpg')))\n        self.transforms = transforms\n        # self.target_transform = target_transform\n\n    def __len__(self):\n        return len(self.dataset)\n\n    def __getitem__(self, idx):\n        path = self.dataset[idx]\n\n        pil_image = self.pil_loader(path)\n        if self.transforms is not None:\n            pil_image = self.transforms(pil_image)\n\n        return pil_image, path\n\n        # if self.target_transform is not None:\n        #     target = self.target_transform(target)\n\n    def pil_loader(self, path):\n        # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\n        with open(path, 'rb') as f:\n            img = Image.open(f)\n            return img.convert('RGB')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nloaders.py\n\"\"\"\n\nimport torch\nimport torch.utils.data\nfrom torchvision import datasets\n\n#from utils import plot_images\n\n#from samplers import valid_and_train_samplers\n#from dataset import TestImageFolder\n\n\ndef get_train_valid_loader(data_dir: str,\n                           batch_size: int,\n                           data_transforms: dict,\n                           random_state: int,\n                           weighted_sampler: bool,\n                           valid_size: float,\n                           shuffle: bool,\n                           show_sample: bool,\n                           num_workers: int,\n                           pin_memory: bool):\n    \"\"\"\n    Utility function for loading and returning train and valid.\n    If using CUDA, num_workers should be set to 1 and pin_memory to True.\n    Params\n    ------\n    - data_dir: path directory to the dataset.\n    - batch_size: how many samples per batch to load.\n    - data_transforms: whether to apply the data augmentation scheme\n    - random_state: fix seed for reproducibility.\n    - valid_size: percentage split of the training set used for\n      the validation set. Should be a float in the range [0, 1].\n    - shuffle: whether to shuffle the train/validation indices.\n    - show_sample: plot 3x8 sample grid of the dataset.\n    - num_workers: number of subprocesses to use when loading the dataset.\n    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to\n      True if using GPU.\n    - weighted: whether to use a weighted sampler\n    Returns\n    -------\n    - train_loader: training set iterator.\n    - valid_loader: validation set iterator.\n    \"\"\"\n\n    # load the dataset\n    train_dataset = datasets.ImageFolder(data_dir, data_transforms['train'])\n    valid_dataset = datasets.ImageFolder(data_dir, data_transforms['valid'])\n\n    train_sampler, valid_sampler, cls_to_weight = valid_and_train_samplers(\n        train_dataset, weighted_sampler, valid_size, random_state)\n\n    train_dataset.cls_to_weight = cls_to_weight\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=num_workers, \n        pin_memory=pin_memory,\n    )\n    valid_loader = torch.utils.data.DataLoader(\n        valid_dataset, batch_size=batch_size, sampler=valid_sampler, num_workers=num_workers, \n        pin_memory=pin_memory,\n    )\n\n    print(f'Total: {len(train_dataset)}; Train/Valid: {len(train_sampler)}/{len(valid_sampler)}')\n\n    # visualize some images\n    if show_sample:\n        sample_loader = torch.utils.data.DataLoader(\n            train_dataset, batch_size=5*8, sampler=train_sampler, num_workers=num_workers,\n            pin_memory=pin_memory,\n        )\n        data_iter = iter(sample_loader)\n        images, labels = data_iter.next()\n        X = images.numpy().transpose([0, 2, 3, 1])\n        plot_images(X, data_dir, labels)\n\n    return (train_loader, valid_loader)\n\n\ndef get_test_loader(data_dir: str,\n                    batch_size: int,\n                    data_transforms: dict,\n                    num_workers: int,\n                    pin_memory: bool):\n    \"\"\"\n    Params\n    ------\n    - data_dir: path directory to the dataset.\n    - batch_size: how many samples per batch to load.\n    - data_transforms: whether to apply the data augmentation scheme\n    - num_workers: number of subprocesses to use when loading the dataset.\n    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to\n      True if using GPU.\n    Returns\n    -------\n    - test_loader: test set iterator.\n    \"\"\"\n\n    test_dataset = TestImageFolder(data_dir, data_transforms['valid'])\n    test_loader = torch.utils.data.DataLoader(\n        test_dataset, batch_size, shuffle=False, num_workers=num_workers,\n        pin_memory=pin_memory,\n    )\n    return test_loader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Model architectures**"},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nmodels.py\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nfrom torchvision import models\n\n\ndef set_parameter_requires_grad(model, feature_extracting):\n    if feature_extracting:\n        for param in model.parameters():\n            param.requires_grad = False\n\n\ndef initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):\n    # Initialize these variables which will be set in this if statement. Each of these\n    #   variables is model specific.\n    model_ft = None\n    input_size = 0\n\n    if model_name == \"resnet\":\n        \"\"\" Resnet18\n        \"\"\"\n        model_ft = models.resnet18(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"alexnet\":\n        \"\"\" Alexnet\n        \"\"\"\n        model_ft = models.alexnet(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.classifier[6].in_features\n        model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"vgg\":\n        \"\"\" VGG11_bn\n        \"\"\"\n        model_ft = models.vgg11_bn(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.classifier[6].in_features\n        model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"squeezenet\":\n        \"\"\" Squeezenet\n        \"\"\"\n        model_ft = models.squeezenet1_0(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        model_ft.classifier[1] = nn.Conv2d(\n            512, num_classes, kernel_size=(1, 1), stride=(1, 1))\n        model_ft.num_classes = num_classes\n        input_size = 224\n\n    elif model_name == \"densenet\":\n        \"\"\" Densenet\n        \"\"\"\n        model_ft = models.densenet121(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.classifier.in_features\n        model_ft.classifier = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"inception\":\n        \"\"\" Inception v3\n        Be careful, expects (299,299) sized images and has auxiliary output\n        \"\"\"\n        model_ft = models.inception_v3(pretrained=use_pretrained)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        # Handle the auxilary net\n        num_ftrs = model_ft.AuxLogits.fc.in_features\n        model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)\n        # Handle the primary net\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 299\n\n    elif model_name == \"resnext50_32x4d\":\n        \"\"\" resnext50_32x4d\n        \"\"\"\n        model_ft = models.resnext50_32x4d(pretrained=True)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n    \n    elif model_name == \"resnext101_32x48d_wsl\":\n        \"\"\" resnext101_32x48d_wsl\n        \"\"\"\n        # model_ft = models.resnext101_32x8d(pretrained=True)\n        model_ft = torch.hub.load('facebookresearch/WSL-Images', 'resnext101_32x48d_wsl')\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n    \n    elif model_name == \"resnext101_32x16d_wsl\":\n        \"\"\" resnext101_32x16d_wsl\n        \"\"\"\n        # model_ft = models.resnext101_32x8d(pretrained=True)\n        model_ft = torch.hub.load('facebookresearch/WSL-Images', 'resnext101_32x16d_wsl')\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"resnext101_32x32d_wsl\":\n        \"\"\" resnext101_32x32d_wsl\n        \"\"\"\n        model_ft = torch.hub.load(\n            'facebookresearch/WSL-Images', 'resnext101_32x32d_wsl')\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    elif model_name == \"resnext101_32x8d\":\n        \"\"\" resnext101_32x8d\n        \"\"\"\n        model_ft = models.resnext101_32x8d(pretrained=True)\n        set_parameter_requires_grad(model_ft, feature_extract)\n        num_ftrs = model_ft.fc.in_features\n        model_ft.fc = nn.Linear(num_ftrs, num_classes)\n        input_size = 224\n\n    else:\n        print(\"Invalid model name, exiting...\")\n        exit()\n\n    return model_ft, input_size","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Train / Test**"},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\ntrain_model.py\n\"\"\"\n\nfrom tqdm import tqdm\nimport torch\nimport numpy as np\nimport time\nimport os\nimport copy\n\nfrom typing import Union\n\ndef train_model(\n    model, \n    dataloaders, \n    criterion, \n    optimizer, \n    device: torch.device,\n    save_best_model_path: Union[None, str] = None,\n    num_epochs: int = 25, \n    is_inception: bool = False):\n    since = time.time()\n\n    val_acc_history = []\n\n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_acc = 0.0\n\n    for epoch in range(1, num_epochs+1):\n        print()\n        # Each epoch has a training and validation phase\n        for phase in ['train', 'valid']:\n            if phase == 'train':\n                model.train()  # Set model to training mode\n            else:\n                model.eval()   # Set model to evaluate mode\n\n            running_loss = 0.0\n            running_corrects = 0\n            seen_images = 0\n            current_accuracies_pbar = []\n            current_loss_pbar = []\n\n            # Iterate over data.\n            progress_bar = tqdm(dataloaders[phase], desc=f'{phase}: ({epoch}/{num_epochs})')\n            for i, (inputs, labels) in enumerate(progress_bar):\n                inputs = inputs.to(device)\n                labels = labels.to(device)\n\n                # zero the parameter gradients\n                optimizer.zero_grad()\n\n                # forward\n                with torch.set_grad_enabled(phase == 'train'):\n                    # Get model outputs and calculate loss\n                    # Special case for inception because in training it has an auxiliary output. In train\n                    #   mode we calculate the loss by summing the final output and the auxiliary output\n                    #   but in testing we only consider the final output.\n                    if is_inception and phase == 'train':\n                        outputs, aux_outputs = model(inputs)\n                        loss1 = criterion(outputs, labels)\n                        loss2 = criterion(aux_outputs, labels)\n                        loss = loss1 + 0.4*loss2\n                    else:\n                        outputs = model(inputs)\n                        loss = criterion(outputs, labels)\n\n                    _, preds = torch.max(outputs, 1)\n\n                    # add loss to the progress bar\n                    current_accuracies_pbar.append(torch.sum(preds == labels.data).item() / len(inputs))\n                    current_loss_pbar.append(loss.item())\n                    if i % 10 == 0:\n                        desc = f'{phase} ({epoch}/{num_epochs}): '\n                        desc += f'Loss: {np.mean(current_loss_pbar):.5f}; '\n                        desc += f'Acc: {np.mean(current_accuracies_pbar):.5f}'\n\n                        progress_bar.set_description(desc)\n                        current_accuracies_pbar = []\n                        current_loss_pbar = []\n\n                    # backward + optimize only if in training phase\n                    if phase == 'train':\n                        loss.backward()\n                        optimizer.step()\n\n                seen_images += len(labels)\n                # statistics\n                running_loss += loss.item() * inputs.size(0)\n                running_corrects += torch.sum(preds == labels.data)\n\n            epoch_loss = running_loss / seen_images\n            epoch_acc = running_corrects.double() / seen_images\n\n            print(f'{phase} Loss: {epoch_loss:.5f} Acc: {epoch_acc:.5f}')\n\n            # deep copy the model\n            if phase == 'valid' and epoch_acc > best_acc:\n                best_acc = epoch_acc\n                best_epoch = epoch\n                best_loss = epoch_loss\n                best_model_wts = copy.deepcopy(model.state_dict())\n                best_optimizer_wts = copy.deepcopy(optimizer.state_dict())\n            if phase == 'valid':\n                val_acc_history.append(epoch_acc)\n\n    time_elapsed = time.time() - since\n    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')\n    print(f'Best val (@ {best_epoch}/{num_epochs} epoch): acc: {best_acc:4f}; loss: {best_loss:4f}')\n\n    # load best model weights\n    model.load_state_dict(best_model_wts)\n\n    # save best model\n    if save_best_model_path is not None:\n        os.makedirs(os.path.split(save_best_model_path)[0], exist_ok=True)\n        \n        torch.save({\n            'epoch': best_epoch,\n            'model_state_dict': model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'loss': best_loss,\n            'accuracy': best_acc\n        }, save_best_model_path)\n\n    # train extra epochs on \n\n    return model, val_acc_history","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\ntest_model.py\n\"\"\"\n\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\n#from utils import save_predictions\nfrom typing import Dict, Union\n\ndef test_model(\n        model: nn.Module,\n        dataloaders: Dict[str, torch.utils.data.dataloader.DataLoader],\n        device: torch.device,\n        save_pred_path: Union[None, str] = None,\n        is_inception: bool = False\n    ) -> None:\n\n    # Set model to evaluate mode\n    model.eval()\n\n    # Initialize predictions list\n    predictions = []\n\n    # class_to_idx dict (class: idx) and reverse\n    class_to_idx = dataloaders['train'].dataset.class_to_idx\n    idx_to_class = {idx: cls for cls, idx in class_to_idx.items()}\n\n    # Iterate over data.\n    progress_bar = tqdm(dataloaders['test'], desc=f'Test: ')\n\n    for i, (inputs, paths) in enumerate(progress_bar):\n        inputs = inputs.to(device)\n\n        with torch.no_grad():\n            outputs = model(inputs)\n            _, preds = torch.max(outputs, 1)\n        \n        # accumulate the preductions after each batch\n        assert len(preds.shape) == 1\n        predictions.extend(list(zip(paths, preds.tolist())))\n\n    # save predictions\n    if save_pred_path is not None:\n        save_predictions(save_pred_path, predictions, idx_to_class)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nfrom torchvision import transforms\nfrom time import localtime, strftime\nimport os\n\n# from train import train_model\n# from test import test_model\n# from loaders import get_train_valid_loader, get_test_loader\n# from models import initialize_model\n# from utils import workdir_copy\n\n# from transforms import ImgAugTransform\n\n# PyTorch Version:  1.3.1\n# Torchvision Version:  0.4.2\n\ndef main():\n    '''\n    Run as: (python ./part2/main.py 2>&1) | tee /home/hdd/logs/openimg/$(date +'%y%m%d%H%M%S').txt\n    '''\n    # save the experiment time\n    start_time = strftime(\"%y%m%d%H%M%S\", localtime())\n\n    # checks and logs\n    pwd = os.getcwd()\n    #assert os.getcwd().endswith('VehicleRecognition')\n    #assert os.path.exists('./part2/experiments/')\n    os.makedirs('./part2/experiments/', exist_ok=True)\n\n    print(f'Working dir: {pwd}')\n    # fix the random seed\n    seed = 13\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n    # paths to dataset\n    # train_data_dir = '/home/nvme/data/openimg/train/train/'\n    # test_data_dir = '/home/nvme/data/openimg/test/testset/'\n    train_data_dir = '../input/vehicle/train/train/'\n    test_data_dir = '../input/vehicle/train/train/'\n\n    # Number of classes in the dataset\n    num_classes = len(os.listdir(train_data_dir))\n\n    # define the paths\n    save_pred_path = None\n    save_pred_path = f'./part2/experiments/{start_time}.csv'\n    save_best_model_path = f'/home/hdd/logs/openimg/{start_time}/best_model.pt'\n\n    # backup the working directiory\n    workdir_copy(pwd, os.path.split(save_best_model_path)[0])\n\n    # resnet, alexnet, vgg, squeezenet, densenet, inception\n    # 'resnext50_32x4d', 'resnext101_32x8d', 'resnext101_32x48d_wsl', 'resnext101_32x32d_wsl'\n    # 'resnext101_32x16d_wsl\n    model_name = \"resnext101_32x16d_wsl\"\n    # Flag for feature extracting. When False, we finetune the whole model,\n    #   when True we only update the reshaped layer params\n    feature_extract = False\n    # hyper parameters\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #torch.device(\"cuda:0\")\n    valid_size = 0.10\n    \n    if model_name.startswith('resnext'):\n        lr = 5e-7\n        # batch_size = 32\n        batch_size = 8\n    elif model_name.startswith('densenet'):\n        lr = 1e-5\n        batch_size = 32\n    else:\n        lr = 1e-4\n        batch_size = 64\n        \n    num_workers = 16\n    pin_memory = True\n    weighted_train_sampler = False\n    weighted_loss = False\n    num_epochs = 20\n\n    # preventing pytorch from allocating some memory on default GPU (0)\n    #torch.cuda.set_device(device)\n\n    # Initialize the model for this run\n    model_ft, input_size = initialize_model(\n        model_name, num_classes, feature_extract, use_pretrained=True)\n\n    # Data augmentation and normalization for training\n    # Just normalization for validation\n    means = [0.485, 0.456, 0.406]\n    stds = [0.229, 0.224, 0.225]\n    data_transforms = {\n        'train': transforms.Compose([\n            transforms.Resize(input_size),\n            transforms.RandomCrop(input_size),\n            transforms.RandomHorizontalFlip(),\n            ImgAugTransform(input_size, 0.25),\n            transforms.ToPILImage(),\n            transforms.ToTensor(),\n            transforms.Normalize(means, stds),\n        ]),\n        'valid': transforms.Compose([\n            transforms.Resize(input_size),\n            transforms.CenterCrop(input_size),\n            transforms.ToTensor(),\n            transforms.Normalize(means, stds)\n        ]),\n    }\n\n    train_loader, valid_loader = get_train_valid_loader(\n        train_data_dir, batch_size, data_transforms, seed, weighted_train_sampler,\n        valid_size=valid_size, shuffle=True, show_sample=True, num_workers=num_workers,\n        pin_memory=pin_memory\n    )\n\n    test_loader = get_test_loader(\n        test_data_dir, batch_size, data_transforms, num_workers=num_workers, pin_memory=pin_memory)\n\n    dataloaders_dict = {\n        'train': train_loader,\n        'valid': valid_loader,\n        'test': test_loader\n    }\n\n    # Send the model to GPU\n    model_ft = model_ft.to(device)\n\n    # Gather the parameters to be optimized/updated in this run. If we are\n    #  finetuning we will be updating all parameters. However, if we are\n    #  doing feature extract method, we will only update the parameters\n    #  that we have just initialized, i.e. the parameters with requires_grad\n    #  is True.\n    params_to_update = model_ft.parameters()\n    print(\"Params to learn:\")\n    if feature_extract:\n        params_to_update = []\n        for name, param in model_ft.named_parameters():\n            if param.requires_grad == True:\n                params_to_update.append(param)\n                print(\"\\t\", name)\n    else:\n        for name, param in model_ft.named_parameters():\n            if param.requires_grad == True:\n                print(\"\\t\", name)\n\n    # Observe that all parameters are being optimized\n    optimizer_ft = optim.Adam(params_to_update, lr=lr)\n\n    # Setup the loss fxn\n    if weighted_loss:\n        print('Weighted Loss')\n        # {0: 0.010101, 1: 0.006622, 2: 0.0008244, 3: 0.00015335, 4: 0.0006253, 5: 0.00019665,\n        # 6: 0.02631, 7: 0.00403, 8: 0.001996, 9: 0.01818, 10: 0.0004466, 11: 0.008771, 12: 0.01087,\n        # 13: 0.006493, 14: 0.0017, 15: 0.000656, 16: 0.001200}\n        cls_to_weight = train_loader.dataset.cls_to_weight\n        weights = torch.FloatTensor([cls_to_weight[c] for c in range(num_classes)]).to(device)\n    else:\n        weights = torch.FloatTensor([1.0 for c in range(num_classes)]).to(device)\n\n    criterion = nn.CrossEntropyLoss(weights)\n\n    # print some things here so it will be seen in terminal for longer time\n    print(f'Timestep: {start_time}')\n    print(f'using model: {model_name}')\n    print(f'Using optimizer: {optimizer_ft}')\n    print(f'Device {device}')\n    print(f'Batchsize: {batch_size}')\n    print(f'Transforms: {data_transforms}')\n\n    # Train and evaluate\n    model_ft, hist = train_model(\n        model_ft, dataloaders_dict, criterion, optimizer_ft, device, save_best_model_path,\n        num_epochs=num_epochs, is_inception=(model_name == \"inception\")\n    )\n\n    # do test inference\n    if save_pred_path is not None:\n        test_model(model_ft, dataloaders_dict, device, save_pred_path,\n                   is_inception=(model_name == \"inception\"))\n\n\nif __name__ == \"__main__\":\n    # parser = argparse.ArgumentParser(description='Training experiment.')\n    # parser.add_argument('--epochs', type=int, default=10)\n    # args = parser.parse_args()\n    # print(args)\n    main()","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}