{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Huge thanks to the starter notebook here! (https://www.kaggle.com/ttahara/training-birdsong-baseline-resnest50-fast). \n### This notebook draws extensively upon the kernel for data IO & hyperparamters.\n### In using it for developing baseline I found it necessary to use mixed precision, allow for different models, and customize my own training loop. \n### So here is this kernel for anyone who wants to start with more custom control over training","execution_count":null},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!pip install ../input/pytorch-pfn-extras/pytorch-pfn-extras-0.2.1/ > /dev/null\n!pip install ../input/resnest50-fast-package/resnest-0.0.6b20200701/resnest/ > /dev/null\n!pip install efficientnet_pytorch > /dev/null\n!pip install timm > /dev/null\n!pip install pretrainedmodels > /dev/null","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.nn.functional as F\ntorch.backends.cudnn.benchmark = True\ntorch.multiprocessing.set_sharing_strategy('file_system')\n\nimport os\nimport cv2\nimport math \nimport timm \nimport shutil\n# import neptune\nimport argparse \nimport numpy as np\nimport pandas as pd\nimport pretrainedmodels\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nfrom PIL import Image\nfrom time import time\nfrom tqdm.autonotebook import tqdm\nfrom warnings import filterwarnings\n# from sync_batchnorm import convert_model\nfrom efficientnet_pytorch import EfficientNet\nfrom torch.utils.tensorboard import SummaryWriter\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset\n\nfrom pathlib import Path\nimport librosa\nimport audioread\nimport soundfile as sf\nimport resnest.torch as resnest_torch\n\npd.options.display.max_rows = 500\npd.options.display.max_columns = 500","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%bash\ngit clone https://github.com/NVIDIA/apex\ncd apex\npip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"args = {\n    'exp': 'starter_development',\n    'model': 'resnest101',\n    'batch_size': 32,\n    'fold': 0,\n    'epochs': 40,\n    'cycle_epochs': 10,\n    'img_size': 224,\n    'margin': 0.2,\n    's': 16.,\n    'fp16': 1,\n    'accum_steps': 1,\n    'pretrained': 'none',\n    'cache_folder': None,\n    'clear_cache': 0,\n    'lr': 6e-4,\n    'n_cpus': 4,\n    'verbose': True,\n    'nowarnings': True\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Warnings and preliminary settings\n# Initialize the neptune experiment\n\n# Logging, caching, and importing\nargs['fp16'] = (args['fp16'] == 1)\nif args['nowarnings']:\n    filterwarnings('ignore')\nif args['fp16']:\n    from apex import amp\nif not os.path.isdir('experiments'):\n    os.mkdir('experiments')\n# if not os.path.isdir('data/cache'):\n#     os.mkdir('data/cache')\nlogging_folder = os.path.join('experiments', args['exp'])\nif os.path.isdir(logging_folder):\n    shutil.rmtree(logging_folder)\nos.mkdir(logging_folder)\n'''if args['cache_folder'] is None:\n    cache_folder = os.path.join('data/cache', str(args['img_size']))\nelse:\n    cache_folder = args['cache_folder']\nif args['clear_cache'] != 0 and os.path.exists(cache_folder):\n    shutil.rmtree(cache_folder)\n    print('Cache folder emptied')\nif not os.path.exists(cache_folder):\n    os.mkdir(cache_folder)\n    print('First time caching')\nprint('Caching preprocessed data@', cache_folder)'''\nwriter = SummaryWriter(logging_folder)\nprint('For visualization, run:')\nprint('tensorboard --logdir='+logging_folder)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"root = '/kaggle/input'\nraw_root = os.path.join(root, \"birdsong-recognition\")\n# TRAIN_AUDIO_DIR = RAW_DATA / \"train_audio\"\ntrain_raw_dir = os.path.join(raw_root, 'train_audio')\ntrain_resampled_dirs = [\n  os.path.join(root, \"birdsong-resampled-train-audio-{:0>2}\".format(i))  for i in range(5)\n]\ntest_raw_dir = os.path.join(raw_root, 'test_audio')\n\n# set(train_df.columns) - set(train_orig_df.columns)\n# Unique columns: {'resampled_channels', 'resampled_filename', 'resampled_sampling_rate'}\n# Resampling rate is invariably 32000, while resample channels is invariably 1 (mono)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join(train_resampled_dirs[0], \"train_mod.csv\"))\ntrain_orig_df = pd.read_csv(os.path.join(raw_root, 'train.csv'))\n\nif not os.path.exists(test_raw_dir):\n    print('No test found (should be the case)')\n    test_raw_dir = os.path.join(root, 'birdcall-check', 'test_audio')\n    test_df = pd.read_csv(os.path.join(root, 'birdcall-check', 'test.csv'))\nelse:\n    test_df = pd.read_csv(os.path.join(test_raw_dir, 'test.csv'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Model portion\n\n# Swish utilities\nfrom efficientnet_pytorch.utils import Swish, MemoryEfficientSwish\nfrom timm.models.layers.activations import Mish\nfrom timm.models.layers.activations import Swish as Swish_timm \n\n### Transform model to (memory efficient) swish ###\ndef to_swish(model):\n    for child_name, child in model.named_children():\n        if isinstance(child, (nn.ReLU, Swish, MemoryEfficientSwish, Swish_timm)):\n            setattr(model, child_name, MemoryEfficientSwish())\n        else:\n            to_swish(child)\n            \ndef to_traceswish(model):\n    for child_name, child in model.named_children():\n        if isinstance(child, MemoryEfficientSwish):\n            setattr(model, child_name, Swish())\n        else:\n            to_traceswish(child)\n            \nclass normalized_Linear(nn.Module):\n    def __init__(self, in_features, out_features, init_s=args['s']):\n        super().__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))\n        self.s = init_s\n        # Important! Initialize it!\n        nn.init.xavier_uniform_(self.weight)\n\n    def forward(self, input):\n        input = F.normalize(input)\n        cos = F.linear(input, F.normalize(self.weight))\n        return self.s*cos\n\nclass add_tail(nn.Module):\n    def __init__(self, backbone, num_features):\n        super().__init__()\n        self.backbone = backbone\n        self.fc = normalized_Linear(num_features, 264)\n    \n    def forward(self, x):\n        return self.fc(self.backbone(x))\n\n# Plain model loading\ndef get_model(args):\n    total_names = ['efficientnet-b'+str(i) for i in range(8)] + pretrainedmodels.model_names\\\n                    + timm.list_models() + ['resnest'+str(depth) for depth in [50, 101, 200, 269]]\n    if not args['model'] in total_names:\n        print('Nope! Available models are:', total_names)\n    # Load from pretrained first\n    if 'efficientnet' in args['model']:\n        try:\n            backbone = EfficientNet.from_pretrained(args['model'], 10)\n        except:\n            print('efficientnet-bx x~[0-7] please')\n            raise NotImplementedError\n        num_features = backbone._fc.weight.shape[1]\n        backbone._fc = nn.Sequential()\n    elif args['model'] in pretrainedmodels.model_names:\n        backbone = pretrainedmodels.__dict__[args['model']](pretrained='imagenet')\n        num_features = backbone.last_linear.weight.shape[1]\n        backbone.last_linear = nn.Sequential()\n    elif 'resnest' in args['model']:\n        backbone = getattr(resnest_torch, args[\"model\"])(pretrained=True)\n        num_features = backbone.fc.weight.shape[1]\n        backbone.fc = nn.Sequential()\n    else:\n        backbone = timm.create_model(args['model'], pretrained=True)\n        for child_name, child in list(backbone.named_children())[::-1]:\n            if isinstance(child, nn.Linear):\n                num_features = child.weight.shape[1]\n                setattr(backbone, child_name, nn.Sequential())\n                break\n    model = add_tail(backbone, num_features)\n    to_swish(model)\n    return model\n\nmodel = get_model(args).cuda()","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"BIRD_CODE = {\n    'aldfly': 0, 'ameavo': 1, 'amebit': 2, 'amecro': 3, 'amegfi': 4,\n    'amekes': 5, 'amepip': 6, 'amered': 7, 'amerob': 8, 'amewig': 9,\n    'amewoo': 10, 'amtspa': 11, 'annhum': 12, 'astfly': 13, 'baisan': 14,\n    'baleag': 15, 'balori': 16, 'banswa': 17, 'barswa': 18, 'bawwar': 19,\n    'belkin1': 20, 'belspa2': 21, 'bewwre': 22, 'bkbcuc': 23, 'bkbmag1': 24,\n    'bkbwar': 25, 'bkcchi': 26, 'bkchum': 27, 'bkhgro': 28, 'bkpwar': 29,\n    'bktspa': 30, 'blkpho': 31, 'blugrb1': 32, 'blujay': 33, 'bnhcow': 34,\n    'boboli': 35, 'bongul': 36, 'brdowl': 37, 'brebla': 38, 'brespa': 39,\n    'brncre': 40, 'brnthr': 41, 'brthum': 42, 'brwhaw': 43, 'btbwar': 44,\n    'btnwar': 45, 'btywar': 46, 'buffle': 47, 'buggna': 48, 'buhvir': 49,\n    'bulori': 50, 'bushti': 51, 'buwtea': 52, 'buwwar': 53, 'cacwre': 54,\n    'calgul': 55, 'calqua': 56, 'camwar': 57, 'cangoo': 58, 'canwar': 59,\n    'canwre': 60, 'carwre': 61, 'casfin': 62, 'caster1': 63, 'casvir': 64,\n    'cedwax': 65, 'chispa': 66, 'chiswi': 67, 'chswar': 68, 'chukar': 69,\n    'clanut': 70, 'cliswa': 71, 'comgol': 72, 'comgra': 73, 'comloo': 74,\n    'commer': 75, 'comnig': 76, 'comrav': 77, 'comred': 78, 'comter': 79,\n    'comyel': 80, 'coohaw': 81, 'coshum': 82, 'cowscj1': 83, 'daejun': 84,\n    'doccor': 85, 'dowwoo': 86, 'dusfly': 87, 'eargre': 88, 'easblu': 89,\n    'easkin': 90, 'easmea': 91, 'easpho': 92, 'eastow': 93, 'eawpew': 94,\n    'eucdov': 95, 'eursta': 96, 'evegro': 97, 'fiespa': 98, 'fiscro': 99,\n    'foxspa': 100, 'gadwal': 101, 'gcrfin': 102, 'gnttow': 103, 'gnwtea': 104,\n    'gockin': 105, 'gocspa': 106, 'goleag': 107, 'grbher3': 108, 'grcfly': 109,\n    'greegr': 110, 'greroa': 111, 'greyel': 112, 'grhowl': 113, 'grnher': 114,\n    'grtgra': 115, 'grycat': 116, 'gryfly': 117, 'haiwoo': 118, 'hamfly': 119,\n    'hergul': 120, 'herthr': 121, 'hoomer': 122, 'hoowar': 123, 'horgre': 124,\n    'horlar': 125, 'houfin': 126, 'houspa': 127, 'houwre': 128, 'indbun': 129,\n    'juntit1': 130, 'killde': 131, 'labwoo': 132, 'larspa': 133, 'lazbun': 134,\n    'leabit': 135, 'leafly': 136, 'leasan': 137, 'lecthr': 138, 'lesgol': 139,\n    'lesnig': 140, 'lesyel': 141, 'lewwoo': 142, 'linspa': 143, 'lobcur': 144,\n    'lobdow': 145, 'logshr': 146, 'lotduc': 147, 'louwat': 148, 'macwar': 149,\n    'magwar': 150, 'mallar3': 151, 'marwre': 152, 'merlin': 153, 'moublu': 154,\n    'mouchi': 155, 'moudov': 156, 'norcar': 157, 'norfli': 158, 'norhar2': 159,\n    'normoc': 160, 'norpar': 161, 'norpin': 162, 'norsho': 163, 'norwat': 164,\n    'nrwswa': 165, 'nutwoo': 166, 'olsfly': 167, 'orcwar': 168, 'osprey': 169,\n    'ovenbi1': 170, 'palwar': 171, 'pasfly': 172, 'pecsan': 173, 'perfal': 174,\n    'phaino': 175, 'pibgre': 176, 'pilwoo': 177, 'pingro': 178, 'pinjay': 179,\n    'pinsis': 180, 'pinwar': 181, 'plsvir': 182, 'prawar': 183, 'purfin': 184,\n    'pygnut': 185, 'rebmer': 186, 'rebnut': 187, 'rebsap': 188, 'rebwoo': 189,\n    'redcro': 190, 'redhea': 191, 'reevir1': 192, 'renpha': 193, 'reshaw': 194,\n    'rethaw': 195, 'rewbla': 196, 'ribgul': 197, 'rinduc': 198, 'robgro': 199,\n    'rocpig': 200, 'rocwre': 201, 'rthhum': 202, 'ruckin': 203, 'rudduc': 204,\n    'rufgro': 205, 'rufhum': 206, 'rusbla': 207, 'sagspa1': 208, 'sagthr': 209,\n    'savspa': 210, 'saypho': 211, 'scatan': 212, 'scoori': 213, 'semplo': 214,\n    'semsan': 215, 'sheowl': 216, 'shshaw': 217, 'snobun': 218, 'snogoo': 219,\n    'solsan': 220, 'sonspa': 221, 'sora': 222, 'sposan': 223, 'spotow': 224,\n    'stejay': 225, 'swahaw': 226, 'swaspa': 227, 'swathr': 228, 'treswa': 229,\n    'truswa': 230, 'tuftit': 231, 'tunswa': 232, 'veery': 233, 'vesspa': 234,\n    'vigswa': 235, 'warvir': 236, 'wesblu': 237, 'wesgre': 238, 'weskin': 239,\n    'wesmea': 240, 'wessan': 241, 'westan': 242, 'wewpew': 243, 'whbnut': 244,\n    'whcspa': 245, 'whfibi': 246, 'whtspa': 247, 'whtswi': 248, 'wilfly': 249,\n    'wilsni1': 250, 'wiltur': 251, 'winwre3': 252, 'wlswar': 253, 'wooduc': 254,\n    'wooscj2': 255, 'woothr': 256, 'y00475': 257, 'yebfly': 258, 'yebsap': 259,\n    'yehbla': 260, 'yelwar': 261, 'yerwar': 262, 'yetvir': 263\n}\n# Beautiful inversion\nINV_BIRD_CODE = {v: k for k, v in BIRD_CODE.items()}","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# Number of seconds for each crop\nPERIOD = 5\n\n# Might be worth some looking into\ndef mono_to_color(\n    X, mean=None, std=None,\n    norm_max=None, norm_min=None, eps=1e-6\n):\n    # Stack to three channels?\n    X = np.stack([X, X, X], axis=-1)\n\n    # Standardize EACH SAMPLE by their minimum, maximum, and renormalize back to 255?\n    mean = mean or X.mean()\n    X = X - mean\n    std = std or X.std()\n    Xstd = X / (std + eps)\n    _min, _max = Xstd.min(), Xstd.max()\n    norm_max = norm_max or _max\n    norm_min = norm_min or _min\n    if (_max - _min) > eps:\n        # Normalize to [0, 255]\n        V = Xstd\n        V[V < norm_min] = norm_min\n        V[V > norm_max] = norm_max\n        V = 255 * (V - norm_min) / (norm_max - norm_min)\n        V = V.astype(np.uint8)\n    else:\n        # Just zero\n        V = np.zeros_like(Xstd, dtype=np.uint8)\n    return V\n\nclass SpectrogramDataset(Dataset):\n    def __init__(\n        self, file_list, img_size=224,\n        waveform_transforms=None, spectrogram_transforms=None, melspectrogram_parameters={}\n    ):\n        self.file_list = file_list  # list of list: [file_path, ebird_code]\n        self.img_size = img_size\n        self.waveform_transforms = waveform_transforms\n        self.spectrogram_transforms = spectrogram_transforms\n        self.melspectrogram_parameters = melspectrogram_parameters\n\n    def __len__(self):\n        return len(self.file_list)\n\n    def __getitem__(self, idx: int):\n        wav_path, ebird_code = self.file_list[idx]\n        # Read what..? y is signal; sr is sample_rate\n        signal, sample_rate = sf.read(wav_path)\n\n        if self.waveform_transforms:\n            signal = self.waveform_transforms(signal)\n        else:\n            len_signal = len(signal)\n            effective_length = sample_rate * PERIOD\n            # If less than [PERIOD] seconds: pad...zeros??? \n            # Reasonable, because centered@zero\n            if len_signal < effective_length:\n                padded_signal = np.zeros(effective_length, dtype=signal.dtype)\n                start = np.random.randint(effective_length - len_signal)\n                padded_signal[start:start + len_signal] = signal\n                signal = padded_signal.astype(np.float32)\n            # Else: Random crop from the whole file???\n            elif len_signal > effective_length:\n                start = np.random.randint(len_signal - effective_length)\n                signal = signal[start:start + effective_length].astype(np.float32)\n            else:\n                signal = signal.astype(np.float32)\n\n        melspec = librosa.feature.melspectrogram(signal, sr=sample_rate,\\\n                                                 **self.melspectrogram_parameters)\n        # Huh? Take log of Amplitude squared (power)?\n        melspec = librosa.power_to_db(melspec).astype(np.float32)\n\n        if self.spectrogram_transforms:\n            melspec = self.spectrogram_transforms(melspec)\n        \n        # Hmmm. This might be worth some looking into\n        image = mono_to_color(melspec)\n        height, width, _ = image.shape\n        # Simple resize of image\n        image = cv2.resize(image, (int(width * self.img_size / height), self.img_size))\n        # Transpose the axis? \n        image = np.moveaxis(image, 2, 0)\n        image = (image / 255.0).astype(np.float32)\n\n        # labels = np.zeros(len(BIRD_CODE), dtype=\"i\")\n        labels = np.zeros(len(BIRD_CODE), dtype=\"f\")\n        labels[BIRD_CODE[ebird_code]] = 1\n\n        return {'input':image, 'label':labels}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tmp_list = []\nfor audio_dir in train_resampled_dirs:\n    if not os.path.exists(audio_dir):\n        continue\n    for ebird_name in os.listdir(audio_dir):\n        ebird_dir = os.path.join(audio_dir, ebird_name)\n        if os.path.isdir(ebird_dir):\n            for wav_name in os.listdir(ebird_dir):\n                wav_path = os.path.join(ebird_dir, wav_name)\n                tmp_list.append([ebird_name, wav_name, wav_path])\n            \ntrain_wav_path_exist = pd.DataFrame(\n    tmp_list, columns=[\"ebird_code\", \"resampled_filename\", \"file_path\"])\n\ntrain_full_df = pd.merge(\n    train_df, train_wav_path_exist, on=[\"ebird_code\", \"resampled_filename\"], how=\"inner\")\n\nprint(train_df.shape)\nprint(train_wav_path_exist.shape)\nprint(train_full_df.shape)\n# I don't understand\ntrain_full_df['ebird_code'].values\n\n# Create splits\nskf = StratifiedKFold(n_splits=5, random_state=2020, shuffle=True)\n\ntrain_full_df[\"fold\"] = -1\nfor fold_id, (train_index, val_index) in enumerate(skf.split(train_full_df, train_full_df[\"ebird_code\"])):\n    train_full_df.iloc[val_index, -1] = fold_id\n    \nuse_fold = args['fold']\ntrain_file_list = train_full_df.query(\"fold != @use_fold\")[[\"file_path\", \"ebird_code\"]].values.tolist()\nval_file_list = train_full_df.query(\"fold == @use_fold\")[[\"file_path\", \"ebird_code\"]].values.tolist()\nprint(\"[fold {}] train: {}, val: {}\".format(use_fold, len(train_file_list), len(val_file_list)))\n    \n# # check the propotion\nfold_proportion = pd.pivot_table(train_full_df, index=\"ebird_code\", columns=\"fold\", values=\"xc_id\", aggfunc=len)\n# Check that stratified is working well\nfold_proportion","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get specific loaders\nmelspectrogram_parameters = {\n    'n_mels': 128,\n    'fmin': 20,\n    'fmax': 16000\n}\ntrainset = SpectrogramDataset(train_file_list, img_size=args['img_size'],\\\n                                      melspectrogram_parameters=melspectrogram_parameters)\nvalset = SpectrogramDataset(val_file_list, img_size=args['img_size'],\\\n                                      melspectrogram_parameters=melspectrogram_parameters)\n\ntrainloader = DataLoader(trainset, pin_memory=False, shuffle=True,\\\n                         batch_size=args['batch_size'], num_workers=args['n_cpus'])\nvalloader = DataLoader(valset, pin_memory=False, shuffle=False,\\\n                         batch_size=args['batch_size'], num_workers=args['n_cpus'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Optimizer and scheduler\nop = torch.optim.AdamW(model.parameters(), lr=args['lr'], weight_decay=5e-4)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(op, max_lr=args['lr'],     total_steps=args['epochs']*len(trainloader))\n# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(op, len(trainloader)*args['cycle_epochs'])\nif args['fp16']:\n    model, op = amp.initialize(model, op, opt_level='O2', verbosity=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"'''Criteria'''\nimport math\nfrom sklearn.metrics import roc_auc_score, roc_curve, auc\n    \n# Assumes both pytorch tensors\ndef metric(preds, labels):\n    # Just a crude mock at the moment\n    return -criterion(preds, labels).item()\n\n# Here! With effective weighting\ndef criterion(preds, labels, reduction='mean'):\n    labels = labels.type_as(preds)\n    return F.binary_cross_entropy_with_logits(preds, labels, reduction=reduction)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"'''training loop'''\nfrom collections import deque\n\nglobal_steps = 0\nckpts_Q = deque(maxlen=5)\nval_losses, val_metrics = [], []\nfor epoch in range(args['epochs']):\n    init_time = time()\n\n    i1 = trainloader\n    if args['verbose'] == 1:\n        i1 = tqdm(i1)\n\n    # Training \n    model.train()\n    train_losses, train_metrics = [], []\n    train_preds, train_labels = [], []\n    for batch in i1:\n        x, labels = batch['input'].cuda(), batch['label'].cuda()\n        preds = model(x)\n        loss = criterion(preds, labels)\n        \n        if args['fp16']:\n            with amp.scale_loss((loss / args['accum_steps']), op) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            (loss / args['accum_steps']).backward()\n\n        if (global_steps + 1) % args['accum_steps'] == 0:\n            op.step()\n            op.zero_grad()\n        \n        global_steps += 1\n        if global_steps % 1 == 0:\n            # Logging\n            train_preds.append(preds.detach().cpu())\n            train_labels.append(labels.cpu())\n            loss_value = loss.item()\n            lr_value = np.max([group['lr'] for group in op.param_groups])\n            metric_value =  metric(torch.cat(train_preds), torch.cat(train_labels))\n            writer.add_scalar('lr', lr_value, global_steps)\n            train_losses.append(loss_value)\n            train_metrics.append(metric_value)\n            if args['verbose']:\n                i1.set_postfix({'train_loss':round(np.mean(train_losses), 3), 'train_metric':round(train_metrics[-1], 3), 'lr':lr_value})\n        scheduler.step()\n    train_loss_value, train_metric_value = np.mean(train_losses), np.mean(train_metrics)\n    \n    model.eval()\n    with torch.no_grad():\n        l, p = [], []\n        model.eval()\n        i1 = valloader\n        if args['verbose']:\n            i1 = tqdm(i1)\n        count = 0\n        for batch in i1:\n            count += 1\n            x, labels = batch['input'].cuda(), batch['label']\n            preds = model(x)\n            l.append(labels)\n            p.append(preds.cpu())\n\n        labels, preds = torch.cat(l), torch.cat(p)\n        val_loss_value, val_metric_value = criterion(preds, labels).item(), metric(preds, labels)\n\n        print(f\"{epoch+1} loss:{round(train_loss_value, 4)} metric:{round(train_metric_value, 4)} val_loss:{round(val_loss_value, 4)} val_metric:{round(val_metric_value, 4)} time:{round(time()-init_time, 4)}\")\n        init_time = time()\n        val_losses.append(val_loss_value)\n        val_metrics.append(val_metric_value)\n        \n        writer.add_scalars('loss', {'train':train_loss_value, 'val':val_loss_value}, global_steps)\n        writer.add_scalars('metric', {'train':train_metric_value, 'val':val_metric_value}, global_steps)\n        \n        if val_metric_value == max(val_metrics):\n            print('New best')\n            # Tracing, a little bit more complicated\n            model_ = get_model(args).cuda()\n            model_.load_state_dict(model.state_dict())\n            to_traceswish(model_)\n            model_.eval()\n            ckpt = torch.jit.trace(model_, x).cpu()\n            ckpts_Q.append(ckpt)\n            for i, ckpt_ in enumerate(ckpts_Q):\n                torch.jit.save(ckpt_, os.path.join(logging_folder, args['exp']+'_'+str(i)+'.pt'))","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}