train_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e31bcc00-6eda-4b05-b26b-53f12137c4c4",
   "metadata": {},
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d4ba0722-90a7-4efd-a3bb-749753998da5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cpu\n"
     ]
    }
   ],
   "source": [
    "# coding=utf-8\n",
    "from __future__ import absolute_import, print_function\n",
    "import os\n",
    "import json\n",
    "from time import time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn import metrics\n",
    "from glob import glob\n",
    "import albumentations as A\n",
    "import cv2\n",
    "from multiprocessing import Process\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import nn\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torchvision import transforms\n",
    "\n",
    "#from DataSet.dataset import get_iwildcam_loader, data_prefetcher\n",
    "#from Utils.train_utils import cross_entropy,focal_loss, get_optimizer\n",
    "#from Utils.train_utils import mixup_data, mixup_criterion\n",
    "from Models.model_factory import create_model\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"0\"\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print('device:', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37a11738-d100-4e51-b8f2-944b31edee3d",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b5b26f1e-09e5-4729-a1ff-8c74bb384106",
   "metadata": {},
   "outputs": [],
   "source": [
    "class cross_entropy(nn.Module):\n",
    "\t\"\"\" Cross entropy that accepts soft targets\"\"\"\n",
    "\n",
    "\tdef __init__(self, size_average=True):\n",
    "\t\tsuper(cross_entropy, self).__init__()\n",
    "\t\tself.size_average = size_average\n",
    "\n",
    "\tdef forward(self, input, target):\n",
    "\t\tlogsoftmax = nn.LogSoftmax()\n",
    "\t\tif self.size_average:\n",
    "\t\t\treturn torch.mean(torch.sum(-target * logsoftmax(input), dim=1))\n",
    "\t\telse:\n",
    "\t\t\treturn torch.sum(torch.sum(-target * logsoftmax(input), dim=1))\n",
    "\n",
    "\n",
    "class focal_loss(nn.Module):\n",
    "\tdef __init__(self, alpha=1., gamma=1.):\n",
    "\t\tsuper(focal_loss, self).__init__()\n",
    "\t\tself.alpha = alpha\n",
    "\t\tself.gamma = gamma\n",
    "\n",
    "\tdef forward(self, inputs, targets, **kwargs):\n",
    "\t\tCE_loss = nn.CrossEntropyLoss(reduction='none')(inputs, targets)\n",
    "\t\tpt = torch.exp(-CE_loss)\n",
    "\t\tF_loss = self.alpha * ((1 - pt) ** self.gamma) * CE_loss\n",
    "\t\treturn F_loss.mean()\n",
    "    \n",
    "def get_optimizer(params, model):\n",
    "\tparam_groups = model.parameters()\n",
    "\n",
    "\tif params['optim'] == 'adam':\n",
    "\t\toptimizer = torch.optim.Adam(param_groups, lr=params['lr'],weight_decay=params['weight_decay'])\n",
    "\telse:\n",
    "\t\toptimizer = torch.optim.SGD(param_groups, lr=params['lr'], momentum=0.9, nesterov=True,weight_decay=params['weight_decay'])\n",
    "\n",
    "\treturn optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "80ac1b50-81ad-44d3-9ab7-a5939185c8a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mixup_data(x, y, alpha=1.0, use_cuda=True):\n",
    "\t'''Returns mixed inputs, pairs of targets, and lambda'''\n",
    "\tif alpha > 0:\n",
    "\t\tlam = np.random.beta(alpha, alpha)\n",
    "\telse:\n",
    "\t\tlam = 1\n",
    "\n",
    "\tbatch_size = x.size()[0]\n",
    "\tif use_cuda:\n",
    "\t\tindex = torch.randperm(batch_size).cuda()\n",
    "\telse:\n",
    "\t\tindex = torch.randperm(batch_size)\n",
    "\n",
    "\tmixed_x = lam * x + (1 - lam) * x[index, :]\n",
    "\ty_a, y_b = y, y[index]\n",
    "\treturn mixed_x, y_a, y_b, lam\n",
    "\n",
    "\n",
    "def mixup_criterion(criterion, pred, y_a, y_b, lam):\n",
    "\treturn lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37889aa8-8953-4940-9a2a-2968a9d7a533",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2445acc9-2025-4b61-ada7-6e5994bbde87",
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_DATASET = {'CCT': 'iWildCam_2019_CCT', 'iNat': 'iWildCam_2019_iNat_Idaho',\n",
    "                 'IDFG': 'iWildCam_IDFG'}  # _images_small\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6ba2ea7f-354d-4e62-9dd4-e701ac68bd00",
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_augment(p=.5, cut_size=8):\n",
    "    imgaugment = A.Compose([\n",
    "        A.HorizontalFlip(p=0.3),\n",
    "        A.GaussNoise(p=.1),\n",
    "        # A.OneOf([\n",
    "        # A.Blur(blur_limit=3, p=.1),\n",
    "        #\tA.GaussNoise(p=.1),\n",
    "        # ], p=0.2),\n",
    "        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT,\n",
    "                           value=(0, 0, 0), p=.3),\n",
    "        A.RandomBrightnessContrast(p=0.3),\n",
    "        A.HueSaturationValue(\n",
    "            hue_shift_limit=20, sat_shift_limit=20, val_shift_limit=20, p=0.1),\n",
    "        A.Cutout(num_holes=1, max_h_size=cut_size,\n",
    "                 max_w_size=cut_size, p=0.3)\n",
    "    ], p=p)\n",
    "\n",
    "    return imgaugment\n",
    "\n",
    "\n",
    "class iWildCam(Dataset):\n",
    "    def __init__(self, params, mode='train'):\n",
    "        self.mode = mode\n",
    "        self.clahe = params['clahe']\n",
    "        self.gray = params['gray']\n",
    "        if 'train' in mode:\n",
    "            clahe_prob = params['clahe_prob']\n",
    "            gray_prob = params['gray_prob']\n",
    "        elif mode == 'infer':\n",
    "            clahe_prob = 1\n",
    "            gray_prob = 1\n",
    "        else:\n",
    "            clahe_prob = int(params['clahe_prob'] >= 1.0)\n",
    "            gray_prob = int(params['gray_prob'] >= 1.0)\n",
    "        if 'train' in mode:\n",
    "            print('use train augmented mode')\n",
    "            self.augment = params['aug_proba'] > 0\n",
    "            self.label_smooth = params['label_smooth']\n",
    "        else:\n",
    "            self.augment = False\n",
    "            self.label_smooth = False\n",
    "        self.one_hot = params['loss'] != 'focal' if mode != 'infer' else False\n",
    "        self.num_classes = params['num_classes']\n",
    "        self.root = params['data_dir']\n",
    "\n",
    "        mean_values = [0.3297, 0.3819, 0.3637]\n",
    "        std_values = [0.1816, 0.1887, 0.1877]\n",
    "\n",
    "        # mean_values = [0.3045, 0.3625, 0.3575]\n",
    "        # std_values = [0.1801, 0.1870, 0.1920]\n",
    "\n",
    "        self.resize = A.Resize(int(params['height'] * 1.1), int(params['width'] * 1.1), interpolation=cv2.INTER_CUBIC,\n",
    "                               p=1.0)\n",
    "        self.crop = A.RandomCrop(params['height'], params['width'], p=1.0) if 'train' in mode else A.CenterCrop(\n",
    "            params['height'], params['width'], p=1.0)\n",
    "\n",
    "        if self.clahe:\n",
    "            self.imgclahe = A.CLAHE(\n",
    "                clip_limit=2.0, tile_grid_size=(16, 16), p=clahe_prob)\n",
    "        if self.gray:\n",
    "            self.imggray = A.ToGray(p=gray_prob)\n",
    "        if self.augment:\n",
    "            self.imgaugment = image_augment(\n",
    "                params['aug_proba'], params['cut_size'])\n",
    "\n",
    "        self.norm = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=mean_values,\n",
    "                                 std=std_values),\n",
    "        ])\n",
    "        if mode == 'train':\n",
    "            self.file_dir = self.root + 'train_file.csv'  # 'train_file_1.csv'\n",
    "        elif mode == 'dev' or mode == 'val' or mode == 'validation':\n",
    "            self.file_dir = self.root + 'dev_file.csv'\n",
    "        elif mode == 'test' or mode == 'infer':\n",
    "            self.file_dir = self.root + 'test_file.csv'\n",
    "        elif mode == 'train_dev' or mode == 'train_val':\n",
    "            self.file_dir = self.root + 'train_file.csv'\n",
    "            self.file_dir_1 = self.root + 'dev_file.csv'\n",
    "        else:\n",
    "            print('does not exisit!', mode)\n",
    "\n",
    "        data_file = pd.read_csv(self.file_dir)\n",
    "        if mode == 'train':\n",
    "            if not params['CCT']:\n",
    "                data_file = data_file[data_file['dataset'] != 'CCT']\n",
    "            if not params['iNat']:\n",
    "                data_file = data_file[data_file['dataset'] != 'iNat']\n",
    "        if mode == 'train_dev' or mode == 'train_val':\n",
    "            temp = pd.read_csv(self.file_dir_1)\n",
    "            data_file = pd.concat([data_file, temp])\n",
    "\n",
    "        data_file = data_file.mask(\n",
    "            data_file.astype(object).eq('None')).dropna()\n",
    "        data_file['absolute_file_name'] = data_file['file_name'].map(\n",
    "            lambda x: os.path.join(self.root, x))\n",
    "\n",
    "        self.image_files = data_file['absolute_file_name'].values\n",
    "        self.image_ids = data_file['id'].values\n",
    "        print('dataset len:', len(self.image_files))\n",
    "        if mode != 'infer':\n",
    "            self.labels = data_file['category_id'].values\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        id = self.image_ids[index]\n",
    "        image = cv2.imread(self.image_files[index])\n",
    "        if image is not None:\n",
    "            image = self.resize(image=image)['image']\n",
    "            if self.clahe:\n",
    "                image = self.imgclahe(image=image)['image']\n",
    "            if self.augment:\n",
    "                image = self.imgaugment(image=image)['image']\n",
    "            if self.gray:\n",
    "                image = self.imggray(image=image)['image']\n",
    "\n",
    "            # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "            image = self.crop(image=image)['image']\n",
    "            image = self.norm(image)\n",
    "        else:\n",
    "            print(self.image_files[index])\n",
    "\n",
    "        if self.mode != 'infer':\n",
    "            label = self.labels[index]\n",
    "            if self.one_hot:\n",
    "                label = np.eye(self.num_classes)[label]\n",
    "                if self.label_smooth > 0:\n",
    "                    label = (1 - self.label_smooth) * label + \\\n",
    "                        self.label_smooth / self.num_classes\n",
    "        else:\n",
    "            label = 0\n",
    "            if self.one_hot:\n",
    "                label = np.eye(self.num_classes)[label]\n",
    "\n",
    "        return (image, label, id)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_files)\n",
    "\n",
    "    def _category(self):\n",
    "        category2id = {\n",
    "            'empty': 0,\n",
    "            'deer': 1,\n",
    "            'moose': 2,\n",
    "            'squirrel': 3,\n",
    "            'rodent': 4,\n",
    "            'small_mammal': 5,\n",
    "            'elk': 6,\n",
    "            'pronghorn_antelope': 7,\n",
    "            'rabbit': 8,\n",
    "            'bighorn_sheep': 9,\n",
    "            'fox': 10,\n",
    "            'coyote': 11,\n",
    "            'black_bear': 12,\n",
    "            'raccoon': 13,\n",
    "            'skunk': 14,\n",
    "            'wolf': 15,\n",
    "            'bobcat': 16,\n",
    "            'cat': 17,\n",
    "            'dog': 18,\n",
    "            'opossum': 19,\n",
    "            'bison': 20,\n",
    "            'mountain_goat': 21,\n",
    "            'mountain_lion': 22\n",
    "        }\n",
    "\n",
    "        id2category = [\n",
    "            'empty',\n",
    "            'deer',\n",
    "            'moose',\n",
    "            'squirrel',\n",
    "            'rodent',\n",
    "            'small_mammal',\n",
    "            'elk',\n",
    "            'pronghorn_antelope',\n",
    "            'rabbit',\n",
    "            'bighorn_sheep',\n",
    "            'fox',\n",
    "            'coyote',\n",
    "            'black_bear',\n",
    "            'raccoon',\n",
    "            'skunk',\n",
    "            'wolf',\n",
    "            'bobcat',\n",
    "            'cat',\n",
    "            'dog',\n",
    "            'opossum',\n",
    "            'bison',\n",
    "            'mountain_goat',\n",
    "            'mountain_lion',\n",
    "        ]\n",
    "\n",
    "\n",
    "class data_prefetcher():\n",
    "    def __init__(self, loader, label_type='float'):\n",
    "        self.loader = iter(loader)\n",
    "        self.stream = torch.cuda.Stream()\n",
    "        self.label_type = label_type\n",
    "        self.preload()\n",
    "\n",
    "    def preload(self):\n",
    "        try:\n",
    "            self.next_input, self.next_target, self.next_ids = next(\n",
    "                self.loader)\n",
    "        except StopIteration:\n",
    "            self.next_input = None\n",
    "            self.next_target = None\n",
    "            self.next_ids = None\n",
    "            return\n",
    "        with torch.cuda.stream(self.stream):\n",
    "            self.next_input = self.next_input.cuda(non_blocking=True)\n",
    "            self.next_target = self.next_target.cuda(non_blocking=True)\n",
    "            #self.next_ids = self.next_ids.cuda(non_blocking=True)\n",
    "\n",
    "            self.next_input = self.next_input.float()\n",
    "            if self.label_type == 'float':\n",
    "                self.next_target = self.next_target.float()\n",
    "            else:\n",
    "                self.next_target = self.next_target.long()\n",
    "\n",
    "    def next(self):\n",
    "        torch.cuda.current_stream().wait_stream(self.stream)\n",
    "        input = self.next_input\n",
    "        target = self.next_target\n",
    "        ids = self.next_ids\n",
    "        self.preload()\n",
    "        return input, target, ids\n",
    "\n",
    "\n",
    "def get_iwildcam_loader(params, mode='train'):\n",
    "    if mode == 'train' or mode == 'train_val' or mode == 'train_dev':\n",
    "        train_data = iWildCam(params, mode=mode)\n",
    "\n",
    "        train_loader = torch.utils.data.DataLoader(\n",
    "            train_data, batch_size=params['batch_size'], shuffle=True,\n",
    "            num_workers=params['threads'], drop_last=True, pin_memory=True)\n",
    "\n",
    "        dev_data = iWildCam(params, mode='dev')\n",
    "\n",
    "        dev_loader = torch.utils.data.DataLoader(\n",
    "            dev_data, batch_size=params['eval_batch_size'], shuffle=False,\n",
    "            num_workers=params['threads'], drop_last=False, pin_memory=True)\n",
    "        return train_loader, dev_loader\n",
    "    elif mode == 'infer':\n",
    "        test_data = iWildCam(params, mode='infer')\n",
    "\n",
    "        test_loader = torch.utils.data.DataLoader(\n",
    "            test_data, batch_size=params['batch_size'], shuffle=False,\n",
    "            num_workers=params['threads'], drop_last=False, pin_memory=True)\n",
    "        return test_loader\n",
    "    else:\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0088271-32ab-4426-935e-defc9823e012",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "325a0afb-b867-4400-ab19-3b5cc55873ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, data_loader, criterion,use_onehot=True):\n",
    "\ty_pred, y_true, losses=[],[],[]\n",
    "\twith torch.no_grad():\n",
    "\t\tinputs, labels, ids = data_loader.next()\n",
    "\t\twhile inputs is not None:\n",
    "\t\t\tif use_onehot:\n",
    "\t\t\t\ttargets = np.argmax(labels.cpu().detach().numpy(), axis=1)\n",
    "\t\t\telse:\n",
    "\t\t\t\ttargets = labels.cpu().detach().numpy()\n",
    "\t\t\ty_true.extend(targets)\n",
    "\t\t\toutput = model(inputs)\n",
    "\t\t\tloss = criterion(output, labels)\n",
    "\t\t\ty_pred.extend(np.argmax(output.cpu().detach().numpy(), axis=1))\n",
    "\t\t\tlosses.append(loss.cpu().detach().numpy())\n",
    "\n",
    "\t\t\tinputs, labels, ids = data_loader.next()\n",
    "\n",
    "\tacc = metrics.accuracy_score(y_true, y_pred)\n",
    "\tf1 = metrics.f1_score(y_true, y_pred, average='macro')\n",
    "\tloss_val=np.mean(losses)\n",
    "\treturn loss_val, acc, f1\n",
    "\n",
    "def train(params):\n",
    "\n",
    "\tif params['init_model'] is not None:\n",
    "\t\tmodel = torch.load(params['init_model'])\n",
    "\t\tprint('load model', params['init_model'])\n",
    "\telse:\n",
    "\t\tmodel = create_model(\n",
    "\t\t\tparams['Net'],\n",
    "\t\t\tpretrained=params['pretrained'],\n",
    "\t\t\tnum_classes=params['num_classes'],\n",
    "\t\t\tdrop_rate=params['drop_rate'],\n",
    "\t\t\tglobal_pool='avg',\n",
    "\t\t\tbn_tf=False,\n",
    "\t\t\tbn_momentum=0.99,\n",
    "\t\t\tbn_eps=1e-3,\n",
    "\t\t\tcheckpoint_path=params['init_model'],\n",
    "\t\t\tin_chans=3)\n",
    "\n",
    "\toptimizer = get_optimizer(params,model)\n",
    "\tparam_num = sum([p.data.nelement() for p in model.parameters()])\n",
    "\tprint(\"Number of model parameters: {} M\".format(param_num / 1024 / 1024))\n",
    "\tmodel = model.to(device)\n",
    "\tmodel.train()\n",
    "\n",
    "\tif params['lr_schedule']:\n",
    "\t\tscheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=params['lr_decay_epochs'], gamma=0.2)\n",
    "\tif params['loss'] =='ce' or params['loss'] =='cross_entropy':\n",
    "\t\tcriterion = cross_entropy().to(device)\n",
    "\t\tlabel_type = 'float'\n",
    "\telif params['loss'] =='focal':\n",
    "\t\tcriterion = focal_loss(gamma=1.0, alpha=1.0).to(device)\n",
    "\t\tlabel_type='long'\n",
    "\telse:\n",
    "\t\tprint('no exist loss',params['loss'])\n",
    "\ttrain_data_loader, dev_data_loader = get_iwildcam_loader(params,mode=params['mode'])\n",
    "\n",
    "\ttrain_log=[]\n",
    "\tdev_log=[]\n",
    "\tbest_acc, best_f1, best_epoch=0,0,0\n",
    "\tt1 = time()\n",
    "\tprint('begin to train')\n",
    "\tuse_onehot=params['loss']!='focal'\n",
    "\tfor epoch in range(params['epochs']):\n",
    "\t\ttrain_loader = data_prefetcher(train_data_loader,label_type)\n",
    "\t\tinputs, labels, ids = train_loader.next()\n",
    "\t\ti = 0\n",
    "\t\twhile inputs is not None:\n",
    "\t\t\tmixup_now = np.random.random()<params['aug_proba']\n",
    "\t\t\tif params['mixup'] and mixup_now:\n",
    "\t\t\t\tinputs, labels_a, labels_b, lam = mixup_data(inputs, labels,\n",
    "\t\t\t                                               params['mixup_alpha'])\n",
    "\n",
    "\n",
    "\t\t\toptimizer.zero_grad()\n",
    "\t\t\toutput = model(inputs)\n",
    "\t\t\tif params['mixup'] and mixup_now:\n",
    "\t\t\t\tloss = mixup_criterion(criterion, output, labels_a, labels_b, lam)\n",
    "\t\t\telse:\n",
    "\t\t\t\tloss = criterion(output, labels)\n",
    "\t\t\tloss.backward()\n",
    "\t\t\toptimizer.step()\n",
    "\n",
    "\t\t\tif i % params['print_step'] == 0:\n",
    "\t\t\t\tpreds = np.argmax(output.cpu().detach().numpy(), axis=1)\n",
    "\t\t\t\tif use_onehot:\n",
    "\t\t\t\t\ttargets = np.argmax(labels.cpu().detach().numpy(), axis=1)\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\ttargets = labels.cpu().detach().numpy()\n",
    "\t\t\t\tacc = metrics.accuracy_score(targets, preds)\n",
    "\t\t\t\tloss_val = loss.cpu().detach().numpy()\n",
    "\t\t\t\tf1 = metrics.f1_score(targets,preds,average='macro')\n",
    "\t\t\t\ttrain_log.append([epoch,i, loss_val, acc, f1])\n",
    "\t\t\t\tprint(\"epoch: %d, iter: %d, train_loss: %.4f, train_acc: %.4f, train_f1: %.4f, time_cost_per_iter: %.4f s\" % (\n",
    "\t\t\t\tepoch, i, loss_val, acc, f1,(time() - t1)/params['print_step']))\n",
    "\t\t\t\twith open(params['log_dir'] + 'train.tsv', 'a') as f:\n",
    "\t\t\t\t\tf.write('%05d\\t%05d\\t%f\\t%f\\t%f\\n' % (epoch, i, loss_val, acc, f1))\n",
    "\t\t\t\tt1 = time()\n",
    "\n",
    "\t\t\tif (i+1) % params['save_step'] == 0:\n",
    "\t\t\t\tsave_model_path= os.path.join(params['save_dir'], 'model_%d_%d.pkl' % (epoch,i))\n",
    "\t\t\t\ttorch.save(model,save_model_path)\n",
    "\t\t\t\tprint('save model to',save_model_path)\n",
    "\n",
    "\t\t\tif (i+1) % params['eval_step'] == 0:\n",
    "\t\t\t\tt2=time()\n",
    "\t\t\t\tmodel.eval()\n",
    "\t\t\t\tdata_loader = data_prefetcher(dev_data_loader,label_type)\n",
    "\t\t\t\tloss_val, acc, f1 = evaluate(model, data_loader, criterion,use_onehot)\n",
    "\t\t\t\tmodel.train()\n",
    "\t\t\t\tdev_log.append([epoch,i, loss_val, acc, f1])\n",
    "\n",
    "\t\t\t\tif f1 > best_f1:\n",
    "\t\t\t\t\tbest_acc, best_f1, best_epoch = acc, f1, epoch\n",
    "\t\t\t\tprint('[Evaluation] -------------------------------')\n",
    "\t\t\t\tprint(\"epoch: %d, test acc: %.4f, f1-score: %.4f, loss: %.4f, best-f1-score: %.4f, eval_time: %.4f s\" % (\n",
    "\t\t\t\t\tepoch, acc, f1, loss_val, best_f1,time()-t2))\n",
    "\t\t\t\tprint('[Evaluation] -------------------------------')\n",
    "\n",
    "\t\t\t\twith open(params['log_dir'] + 'eval.tsv', 'a') as f:\n",
    "\t\t\t\t\tf.write('%05d\\t%05d\\t%f\\t%f\\t%f\\n' % (epoch, i, loss_val, acc, f1))\n",
    "\n",
    "\t\t\tinputs, labels, ids = train_loader.next()\n",
    "\t\t\ti += 1\n",
    "\n",
    "\t\tif params['lr_schedule']:\n",
    "\t\t\tscheduler.step(epoch)\n",
    "\n",
    "\treturn model\n",
    "\n",
    "def get_params():\n",
    "\tparams = {\n",
    "\t\t'mode':'train_val',\n",
    "\t\t'data_dir': 'data/bbox/cropped_image/', #['data/bbox/cropped_image/','data/']\n",
    "\t\t'CCT':True,\n",
    "\t\t'iNat':True,\n",
    "\t\t'save_dir': 'final_output/output_0/',\n",
    "\t\t'init_model': None,#'output_1/resnet_101_3_3427.pkl',\n",
    "\t\t'Net': 'tf_efficientnet_b0',  # 'resnet','wideresnet','tf_efficientnet_b0'\n",
    "\t\t'pretrained': True,\n",
    "\t\t'drop_rate':0.2,\n",
    "\n",
    "\t\t'batch_size': 32,\n",
    "\t\t'eval_batch_size': 32,\n",
    "\t\t'num_classes': 23,\n",
    "\t\t'epochs': 6,\n",
    "\t\t'print_per_epoch':500,\n",
    "\t\t'eval_per_epoch': 4,\n",
    "\t\t'save_per_epoch': 4,\n",
    "\n",
    "\t\t'loss':'ce',#['ce','focal']\n",
    "\t\t'lr_schedule': True,\n",
    "\t\t'lr': 5e-3,\n",
    "\t\t'weight_decay':1e-6,\n",
    "\t\t'optim': 'adam',\n",
    "\t\t'lr_decay_epochs':[2,4],\n",
    "\n",
    "\t\t'clahe':True,\n",
    "\t\t'clahe_prob': 0.2,\n",
    "\t\t'gray':True,\n",
    "\t\t'gray_prob':0.01,\n",
    "\t\t'aug_proba':0.5,\n",
    "\t\t'cut_size':8,\n",
    "\t\t'label_smooth':0.01,\n",
    "\t\t'mixup':True,\n",
    "\t\t'mixup_alpha':1,\n",
    "\t\t'height':64,#380,#224 resnet, 300\n",
    "\t\t'width':64,\n",
    "\t\t'threads':2,\n",
    "\t}\n",
    "\tparams['log_dir'] = os.path.join(params['save_dir'], 'log/')\n",
    "\tif not os.path.exists(params['save_dir']):\n",
    "\t\tos.mkdir(params['save_dir'])\n",
    "\tif not os.path.exists(params['log_dir']):\n",
    "\t\tos.mkdir(params['log_dir'])\n",
    "\t\twith open(params['log_dir'] + 'eval.tsv', 'a') as f:\n",
    "\t\t\tf.write('Epoch\\tStep\\tLoss\\tAccuracy\\tF1-Score\\n')\n",
    "\t\twith open(params['log_dir'] + 'train.tsv', 'a') as f:\n",
    "\t\t\tf.write('Epoch\\tStep\\tLoss\\tAccuracy\\tF1-Score\\n')\n",
    "\troot = params['data_dir']\n",
    "\tparams['train_data_size'] = len(pd.read_csv(root + 'train_file.csv'))\n",
    "\tparams['dev_data_size'] = len(pd.read_csv(root + 'dev_file.csv'))\n",
    "\tparams['step_per_epoch'] = params['train_data_size'] // params['batch_size']\n",
    "\tparams['print_step'] = max(1,params['step_per_epoch']//params['print_per_epoch'])\n",
    "\tparams['eval_step'] = max(1,params['step_per_epoch']//params['eval_per_epoch'])\n",
    "\tparams['save_step'] = max(1,params['step_per_epoch']//params['save_per_epoch'])\n",
    "\n",
    "\tjson.dump(obj=params, fp=open(params['log_dir'] + 'parameters.json', 'w'))\n",
    "\tprint(params)\n",
    "\n",
    "\treturn params\n",
    "\n",
    "def load_params(save_dir):\n",
    "\tparams_path=save_dir + 'log/parameters.json'\n",
    "\tprint('load params form',params_path)\n",
    "\tparams = json.load(fp=open(params_path, 'r'))\n",
    "\tckpts = glob(save_dir+'*.pkl')\n",
    "\tif len(ckpts)>0:\n",
    "\t\tckpts = sorted(ckpts, key=lambda x: eval(x.split('/')[-1].split('.')[0].split('_')[-1]))\n",
    "\t\tparams['init_model']=ckpts[-1]\n",
    "\tprint(params)\n",
    "\treturn params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "609fe461-5163-40e2-824f-0f094a2cd12f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'mode': 'train_val', 'data_dir': 'data/bbox/cropped_image/', 'CCT': True, 'iNat': True, 'save_dir': 'final_output/output_0/', 'init_model': None, 'Net': 'tf_efficientnet_b0', 'pretrained': True, 'drop_rate': 0.2, 'batch_size': 32, 'eval_batch_size': 32, 'num_classes': 23, 'epochs': 6, 'print_per_epoch': 500, 'eval_per_epoch': 4, 'save_per_epoch': 4, 'loss': 'ce', 'lr_schedule': True, 'lr': 0.005, 'weight_decay': 1e-06, 'optim': 'adam', 'lr_decay_epochs': [2, 4], 'clahe': True, 'clahe_prob': 0.2, 'gray': True, 'gray_prob': 0.01, 'aug_proba': 0.5, 'cut_size': 8, 'label_smooth': 0.01, 'mixup': True, 'mixup_alpha': 1, 'height': 64, 'width': 64, 'threads': 2, 'log_dir': 'final_output/output_0/log/', 'train_data_size': 196332, 'dev_data_size': 6145, 'step_per_epoch': 6135, 'print_step': 12, 'eval_step': 1533, 'save_step': 1533}\n"
     ]
    }
   ],
   "source": [
    "params = get_params()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "032d6121-f7a1-40fc-8ca5-fae236d64543",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cpu\n",
      "{'mode': 'train_val', 'data_dir': 'data/bbox/cropped_image/', 'CCT': True, 'iNat': True, 'save_dir': 'final_output/output_0/', 'init_model': None, 'Net': 'tf_efficientnet_b0', 'pretrained': True, 'drop_rate': 0.2, 'batch_size': 32, 'eval_batch_size': 32, 'num_classes': 23, 'epochs': 6, 'print_per_epoch': 500, 'eval_per_epoch': 4, 'save_per_epoch': 4, 'loss': 'ce', 'lr_schedule': True, 'lr': 0.005, 'weight_decay': 1e-06, 'optim': 'adam', 'lr_decay_epochs': [2, 4], 'clahe': True, 'clahe_prob': 0.2, 'gray': True, 'gray_prob': 0.01, 'aug_proba': 0.5, 'cut_size': 8, 'label_smooth': 0.01, 'mixup': True, 'mixup_alpha': 1, 'height': 64, 'width': 64, 'threads': 2, 'log_dir': 'final_output/output_0/log/', 'train_data_size': 196332, 'dev_data_size': 6145, 'step_per_epoch': 6135, 'print_step': 12, 'eval_step': 1533, 'save_step': 1533}\n",
      "Number of model parameters: 3.8499937057495117 M\n",
      "use train augmented mode\n",
      "dataset len: 202477\n",
      "dataset len: 6145\n",
      "begin to train\n",
      "Traceback (most recent call last):\n",
      "  File \"train_model.py\", line 253, in <module>\n",
      "    main()\n",
      "  File \"train_model.py\", line 249, in main\n",
      "    train(params)\n",
      "  File \"train_model.py\", line 95, in train\n",
      "    train_loader = data_prefetcher(train_data_loader, label_type)\n",
      "  File \"/notebooks/iWildCam_2019_FGVC6/DataSet/dataset.py\", line 204, in __init__\n",
      "    self.stream = torch.cuda.Stream()\n",
      "  File \"/usr/local/lib/python3.8/dist-packages/torch/cuda/streams.py\", line 34, in __new__\n",
      "    with torch.cuda.device(device):\n",
      "  File \"/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py\", line 225, in __enter__\n",
      "    self.prev_idx = torch.cuda.current_device()\n",
      "  File \"/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py\", line 432, in current_device\n",
      "    _lazy_init()\n",
      "  File \"/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py\", line 172, in _lazy_init\n",
      "    torch._C._cuda_init()\n",
      "RuntimeError: No CUDA GPUs are available\n"
     ]
    }
   ],
   "source": [
    "!python3 train_model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4bfc4894-b9c2-4c75-8cd3-dfb120c066dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting albumentations\n",
      "  Downloading albumentations-1.0.3-py3-none-any.whl (98 kB)\n",
      "\u001b[K     |████████████████████████████████| 98 kB 1.0 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting opencv-python-headless>=4.1.1\n",
      "  Downloading opencv_python_headless-4.5.3.56-cp38-cp38-manylinux2014_x86_64.whl (37.1 MB)\n",
      "\u001b[K     |████████████████████████████████| 37.1 MB 8.3 MB/s eta 0:00:011\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.11.1 in /usr/local/lib/python3.8/dist-packages (from albumentations) (1.19.5)\n",
      "Requirement already satisfied: PyYAML in /usr/local/lib/python3.8/dist-packages (from albumentations) (5.4.1)\n",
      "Collecting scikit-image>=0.16.1\n",
      "  Downloading scikit_image-0.18.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (30.2 MB)\n",
      "\u001b[K     |████████████████████████████████| 30.2 MB 11.9 MB/s eta 0:00:01    |██████████████████████▌         | 21.2 MB 9.5 MB/s eta 0:00:01     |███████████████████████▊        | 22.4 MB 9.5 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from albumentations) (1.7.1)\n",
      "Collecting networkx>=2.0\n",
      "  Downloading networkx-2.6.3-py3-none-any.whl (1.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.9 MB 7.5 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: matplotlib!=3.0.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-image>=0.16.1->albumentations) (3.4.3)\n",
      "Collecting PyWavelets>=1.1.1\n",
      "  Downloading PyWavelets-1.1.1-cp38-cp38-manylinux1_x86_64.whl (4.4 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.4 MB 6.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tifffile>=2019.7.26\n",
      "  Downloading tifffile-2021.8.30-py3-none-any.whl (171 kB)\n",
      "\u001b[K     |████████████████████████████████| 171 kB 8.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting imageio>=2.3.0\n",
      "  Downloading imageio-2.9.0-py3-none-any.whl (3.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 3.3 MB 8.9 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: pillow!=7.1.0,!=7.1.1,>=4.3.0 in /usr/local/lib/python3.8/dist-packages (from scikit-image>=0.16.1->albumentations) (8.3.2)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image>=0.16.1->albumentations) (2.4.7)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.8/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image>=0.16.1->albumentations) (2.8.2)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image>=0.16.1->albumentations) (1.3.2)\n",
      "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.8/dist-packages (from matplotlib!=3.0.0,>=2.0.0->scikit-image>=0.16.1->albumentations) (0.10.0)\n",
      "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.7->matplotlib!=3.0.0,>=2.0.0->scikit-image>=0.16.1->albumentations) (1.15.0)\n",
      "Installing collected packages: opencv-python-headless, networkx, PyWavelets, tifffile, imageio, scikit-image, albumentations\n",
      "Successfully installed PyWavelets-1.1.1 albumentations-1.0.3 imageio-2.9.0 networkx-2.6.3 opencv-python-headless-4.5.3.56 scikit-image-0.18.3 tifffile-2021.8.30\n"
     ]
    }
   ],
   "source": [
    "!pip3 install albumentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf722ad6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}