atom-transfomer-good-norm-pytorch-apex.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import apex\n",
    "from apex import amp\n",
    "\n",
    "#%load_ext autoreload\n",
    "#%autoreload 2\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1q\"\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.utils.data\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Parameter\n",
    "\n",
    "def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,\n",
    "                          missing_keys, unexpected_keys, error_msgs):\n",
    "    r\"\"\"Copies parameters and buffers from :attr:`state_dict` into only\n",
    "    this module, but not its descendants. This is called on every submodule\n",
    "    in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this\n",
    "    module in input :attr:`state_dict` is provided as :attr:`local_metadata`.\n",
    "    For state dicts without metadata, :attr:`local_metadata` is empty.\n",
    "    Subclasses can achieve class-specific backward compatible loading using\n",
    "    the version number at `local_metadata.get(\"version\", None)`.\n",
    "\n",
    "    .. note::\n",
    "        :attr:`state_dict` is not the same object as the input\n",
    "        :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So\n",
    "        it can be modified.\n",
    "\n",
    "    Arguments:\n",
    "        state_dict (dict): a dict containing parameters and\n",
    "            persistent buffers.\n",
    "        prefix (str): the prefix for parameters and buffers used in this\n",
    "            module\n",
    "        local_metadata (dict): a dict containing the metadata for this module.\n",
    "            See\n",
    "        strict (bool): whether to strictly enforce that the keys in\n",
    "            :attr:`state_dict` with :attr:`prefix` match the names of\n",
    "            parameters and buffers in this module\n",
    "        missing_keys (list of str): if ``strict=True``, add missing keys to\n",
    "            this list\n",
    "        unexpected_keys (list of str): if ``strict=True``, add unexpected\n",
    "            keys to this list\n",
    "        error_msgs (list of str): error messages should be added to this\n",
    "            list, and will be reported together in\n",
    "            :meth:`~torch.nn.Module.load_state_dict`\n",
    "    \"\"\"\n",
    "    for hook in self._load_state_dict_pre_hooks.values():\n",
    "        hook(state_dict, prefix, local_metadata, strict, missing_keys,\n",
    "             unexpected_keys, error_msgs)\n",
    "\n",
    "    local_name_params = itertools.chain(self._parameters.items(),\n",
    "                                        self._buffers.items())\n",
    "    local_state = {k: v.data for k, v in local_name_params if v is not None}\n",
    "\n",
    "    for name, param in local_state.items():\n",
    "        key = prefix + name\n",
    "        if key in state_dict:\n",
    "            input_param = state_dict[key]\n",
    "\n",
    "            # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+\n",
    "            if len(param.shape) == 0 and len(input_param.shape) == 1:\n",
    "                input_param = input_param[0]\n",
    "\n",
    "            if input_param.shape != param.shape:\n",
    "                # local shape should match the one in checkpoint\n",
    "                error_msgs.append(\n",
    "                    'size mismatch for {}: copying a param with shape {} from checkpoint, '\n",
    "                    'the shape in current model is {}.'.format(\n",
    "                        key, input_param.shape, param.shape))\n",
    "                #if not strict:\n",
    "                #    continue\n",
    "\n",
    "            if isinstance(input_param, Parameter):\n",
    "                # backwards compatibility for serialized parameters\n",
    "                input_param = input_param.data\n",
    "\n",
    "            try:\n",
    "                param.copy_(input_param)\n",
    "            except Exception:\n",
    "                error_msgs.append(\n",
    "                    'While copying the parameter named \"{}\", '\n",
    "                    'whose dimensions in the model are {} and '\n",
    "                    'whose dimensions in the checkpoint are {}.'.format(\n",
    "                        key, param.size(), input_param.size()))\n",
    "                # PG load partially\n",
    "\n",
    "                if len(input_param.size()) == 3:\n",
    "                    error_msgs.append(\n",
    "                        'Partially copying the parameter named \"{}\", '\n",
    "                        'whose dimensions in the model are {} and '\n",
    "                        'whose dimensions in the checkpoint are {}. - trying {}'\n",
    "                        .format(\n",
    "                            key, param.size(), input_param.size(),\n",
    "                            param[:input_param.size()[0], :input_param.size(\n",
    "                            )[1], :input_param.size()[2]].shape))\n",
    "                else:\n",
    "                    error_msgs.append(\n",
    "                        'Partially copying the parameter named \"{}\", '\n",
    "                        'whose dimensions in the model are {} and '\n",
    "                        'whose dimensions in the checkpoint are {}. - trying {}'\n",
    "                        .format(key, param.size(), input_param.size(),\n",
    "                                param[:input_param.size()[0]].shape))\n",
    "\n",
    "                try:\n",
    "                    new_input_param = torch.empty_like(param)\n",
    "                    new_input_param = torch.nn.init.normal_(new_input_param,\n",
    "                                                            mean=input_param.mean(),\n",
    "                                                            std=input_param.std())\n",
    "\n",
    "                    if len(input_param.size()) == 3:\n",
    "                        new_input_param[:input_param.size()[0], :input_param.\n",
    "                                        size()[1], :input_param.size(\n",
    "                                        )[2]] = input_param\n",
    "                    else:\n",
    "                        new_input_param[:input_param.size()[0]] = input_param\n",
    "                    param.copy_(new_input_param)\n",
    "                except Exception as e:\n",
    "                    assert e\n",
    "                    error_msgs.append(\n",
    "                        'Failed to load weights partially {}'.format(e))\n",
    "        elif strict:\n",
    "            missing_keys.append(key)\n",
    "\n",
    "    if strict:\n",
    "        for key in state_dict.keys():\n",
    "            if key.startswith(prefix):\n",
    "                input_name = key[len(prefix):]\n",
    "                input_name = input_name.split(\n",
    "                    '.', 1)[0]  # get the name of param/buffer/child\n",
    "                if input_name not in self._modules and input_name not in local_state:\n",
    "                    unexpected_keys.append(key)\n",
    "\n",
    "def load_state_dict(self, state_dict, strict=True):\n",
    "    r\"\"\"Copies parameters and buffers from :attr:`state_dict` into\n",
    "    this module and its descendants. If :attr:`strict` is ``True``, then\n",
    "    the keys of :attr:`state_dict` must exactly match the keys returned\n",
    "    by this module's :meth:`~torch.nn.Module.state_dict` function.\n",
    "\n",
    "    Arguments:\n",
    "        state_dict (dict): a dict containing parameters and\n",
    "            persistent buffers.\n",
    "        strict (bool, optional): whether to strictly enforce that the keys\n",
    "            in :attr:`state_dict` match the keys returned by this module's\n",
    "            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n",
    "\n",
    "    Returns:\n",
    "        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n",
    "            * **missing_keys** is a list of str containing the missing keys\n",
    "            * **unexpected_keys** is a list of str containing the unexpected keys\n",
    "    \"\"\"\n",
    "    missing_keys = []\n",
    "    unexpected_keys = []\n",
    "    error_msgs = []\n",
    "\n",
    "    # copy state_dict so _load_from_state_dict can modify it\n",
    "    metadata = getattr(state_dict, '_metadata', None)\n",
    "    state_dict = state_dict.copy()\n",
    "    if metadata is not None:\n",
    "        state_dict._metadata = metadata\n",
    "\n",
    "    def load(module, prefix=''):\n",
    "        local_metadata = {} if metadata is None else metadata.get(\n",
    "            prefix[:-1], {})\n",
    "        module._load_from_state_dict(state_dict, prefix, local_metadata, True,\n",
    "                                     missing_keys, unexpected_keys, error_msgs)\n",
    "        for name, child in module._modules.items():\n",
    "            if child is not None:\n",
    "                load(child, prefix + name + '.')\n",
    "                \n",
    "    load(self)\n",
    "\n",
    "    if strict:\n",
    "        if len(unexpected_keys) > 0:\n",
    "            error_msgs.insert(\n",
    "                0, 'Unexpected key(s) in state_dict: {}. '.format(', '.join(\n",
    "                    '\"{}\"'.format(k) for k in unexpected_keys)))\n",
    "        if len(missing_keys) > 0:\n",
    "            error_msgs.insert(\n",
    "                0, 'Missing key(s) in state_dict: {}. '.format(', '.join(\n",
    "                    '\"{}\"'.format(k) for k in missing_keys)))\n",
    "\n",
    "    if strict and len(error_msgs) > 0:\n",
    "        raise RuntimeError(\n",
    "            'Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n",
    "                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn.Module._load_from_state_dict = _load_from_state_dict\n",
    "nn.Module.load_state_dict = load_state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai import *\n",
    "from fastai.basic_train import *\n",
    "from fastai.basic_data import *\n",
    "from fastai.data_block import *\n",
    "from fastai.torch_core import *\n",
    "from fastai.train import *\n",
    "from fastai.callback import *\n",
    "from fastai.callbacks import *\n",
    "from fastai.distributed import *\n",
    "from fastai.layers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# deterministic\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_ext = lambda fname,ext: f'{str(fname)[:-4]}{ext}{str(fname)[-4:]}'\n",
    "\n",
    "def preprocess(fname, type_index=None, ext=''):\n",
    "    t  = pd.read_csv(fname_ext(fname,ext))\n",
    "    s  = pd.read_csv('structures.csv')\n",
    "    \n",
    "    has_y = 'scalar_coupling_constant' in t.columns\n",
    "\n",
    "    if has_y:\n",
    "        # atom-atom level\n",
    "        # molecule_name,atom_index_0,atom_index_1,type,fc,sd,pso,dso\n",
    "        scalar_couplings = pd.read_csv(f'scalar_coupling_contributions{ext}.csv') # fc,sd,pso,dso\n",
    "\n",
    "        # atom level\n",
    "        # molecule_name,atom_index,XX,YX,ZX,XY,YY,ZY,XZ,YZ,ZZ\n",
    "        magnetic_shielding = pd.read_csv('magnetic_shielding_tensors.csv')\n",
    "        # molecule_name,atom_index,mulliken_charge\n",
    "        mulliken_charges = pd.read_csv('mulliken_charges.csv')\n",
    "\n",
    "        # molecule level\n",
    "        # molecule_name,X,Y,Z\n",
    "        dipole_moments = pd.read_csv('dipole_moments.csv')\n",
    "        # molecule_name,potential_energy\n",
    "        potential_energy = pd.read_csv('potential_energy.csv')\n",
    "\n",
    "    t['molecule_index'] = pd.factorize(t['molecule_name'])[0] + t['id'].min()\n",
    "    # make sure we use the same indexes in train/test (test needs to provide type_index)\n",
    "    if type_index is not None:\n",
    "        t['type_idx'] = t['type'].apply(lambda x: type_index.index(x) ) # pd.factorize(pd.concat([pd.Series(type_index),t['type']]))[0][len(type_index):]\n",
    "    else:\n",
    "        t['type_idx'] = pd.factorize(t['type'])[0]\n",
    "\n",
    "    s['atom_idx'] = s['atom'].apply(lambda x: atoms.index(x) )\n",
    "\n",
    "    max_items = len(t.groupby(['molecule_name', 'atom_index_0']))# if has_y else 422550\n",
    "    max_atoms = int(s.atom_index.max() + 1)\n",
    "\n",
    "    if has_y:\n",
    "        contributions = ['fc','sd','pso','dso']\n",
    "        magnetic_tensors = ['XX','YX','ZX','XY','YY','ZY','XZ','YZ','ZZ']\n",
    "        XYZ = ['X','Y','Z']\n",
    "    xyz = ['x', 'y', 'z']\n",
    "    \n",
    "    x_xyz   = np.zeros((max_items,len(xyz),  max_atoms), dtype=np.float32)\n",
    "    x_type  = np.zeros((max_items,1,         max_atoms), dtype=np.int8)\n",
    "    x_ext   = np.zeros((max_items,1,         max_atoms), dtype=np.bool_)\n",
    "    x_atom  = np.empty((max_items,1,         max_atoms), dtype=np.int8)\n",
    "    x_atom[:] = -1\n",
    "\n",
    "    if has_y:\n",
    "        y_scalar   = np.zeros((max_items,len(contributions)   ,max_atoms), dtype=np.float32)\n",
    "        y_magnetic = np.zeros((max_items,len(magnetic_tensors),max_atoms), dtype=np.float32)\n",
    "        y_mulliken = np.zeros((max_items,1                    ,max_atoms), dtype=np.float32)\n",
    "\n",
    "        y_dipole   = np.zeros((max_items,len(XYZ)), dtype=np.float32)\n",
    "        y_potential= np.zeros((max_items,1       ), dtype=np.float32)\n",
    "\n",
    "        y_magnetic[...] = np.nan\n",
    "        y_mulliken[...] = np.nan\n",
    "    else:\n",
    "        xt_ids = np.zeros((max_items, max_atoms), dtype=np.int32)\n",
    "\n",
    "    m = np.zeros((max_items,), dtype=np.int32)\n",
    "    i = j = 0\n",
    "    \n",
    "    for (m_name, m_index) ,m_group in tqdm(t.groupby(['molecule_name', 'molecule_index'])):\n",
    "        ss = s[s.molecule_name==m_name]\n",
    "        n_atoms = len(ss)\n",
    "        if has_y:\n",
    "            magnetic = magnetic_shielding[\n",
    "                    (magnetic_shielding['molecule_name']==m_name)][magnetic_tensors].values.T\n",
    "\n",
    "            mulliken = mulliken_charges[\n",
    "                    (mulliken_charges['molecule_name']==m_name)]['mulliken_charge'].values.T\n",
    "\n",
    "            scs = scalar_couplings[scalar_couplings['molecule_name']==m_name]\n",
    "            \n",
    "            y_dipole[j,:]= dipole_moments[dipole_moments['molecule_name']==m_name][XYZ].values\n",
    "            y_potential[j,:]=potential_energy[\n",
    "                potential_energy['molecule_name']==m_name]['potential_energy'].values\n",
    "        \n",
    "        for a_name,a_group in m_group.groupby('atom_index_0'):\n",
    "            \n",
    "            ref_a = ss[ss['atom_index']==a_name]\n",
    "            \n",
    "            x_xyz[i] = 0.\n",
    "            x_type[i] = -1\n",
    "            x_ext[i] =  True\n",
    "            \n",
    "            x_xyz[i,:,:n_atoms] = (ss[xyz].values-ref_a[xyz].values).T  # xyz \n",
    "            x_type[i,0,a_group['atom_index_1']] = a_group['type_idx']  # type \n",
    "            x_ext[i,0,a_group['atom_index_1']] = a_group['ext']  # ext \n",
    "            x_atom[i,:,:n_atoms] = ss['atom_idx'].T                \n",
    "\n",
    "            if has_y:\n",
    "                y_scalar[i,:,a_group['atom_index_1']] = scs[scs['atom_index_0']==a_name][contributions]\n",
    "                y_magnetic[i,:,:n_atoms] = magnetic\n",
    "                y_mulliken[i,:,:n_atoms] = mulliken\n",
    "            else:\n",
    "                xt_ids[i,a_group['atom_index_1']] = a_group['id']  \n",
    "\n",
    "            m[i] = m_index\n",
    "            i+=1\n",
    "        j += 1\n",
    "    assert i == max_items\n",
    "    print(i,max_items)\n",
    "    if has_y:\n",
    "        return x_xyz,x_type,x_ext,x_atom, m, y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential\n",
    "    else:\n",
    "        return x_xyz,x_type,x_ext,x_atom, m, xt_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define where you want to use original training set '' or extended ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ext = '_ext' # or ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load preprocessed or preprocess and save for later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_fname = Path('train.npz')\n",
    "types = ['1JHC', '2JHH', '1JHN', '2JHN', '2JHC', '3JHH', '3JHC', '3JHN'] \n",
    "atoms = 'CFHNO'\n",
    "\n",
    "try:\n",
    "    npzfile = np.load(fname_ext(train_fname, ext))\n",
    "    x_xyz   = npzfile['x_xyz']\n",
    "    x_type  = npzfile['x_type']\n",
    "    x_ext   = npzfile['x_ext']\n",
    "    x_atom  = npzfile['x_atom']\n",
    "\n",
    "    y_scalar    = npzfile['y_scalar']\n",
    "    y_magnetic  = npzfile['y_magnetic']\n",
    "    y_mulliken  = npzfile['y_mulliken']\n",
    "    y_dipole    = npzfile['y_dipole']\n",
    "    y_potential = npzfile['y_potential']\n",
    "    m = npzfile['m']\n",
    "    max_items, max_atoms = x_xyz.shape[0], x_xyz.shape[-1]\n",
    "except:\n",
    "    x_xyz,x_type,x_ext,x_atom, m, y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential = \\\n",
    "        preprocess(train_fname.with_suffix('.csv'), type_index=types, ext=ext)\n",
    "    np.savez(fname_ext(train_fname, ext), \n",
    "             x_xyz=x_xyz,\n",
    "             x_type=x_type,\n",
    "             x_ext=x_ext,\n",
    "             x_atom=x_atom,\n",
    "             y_scalar=y_scalar,\n",
    "             y_magnetic=y_magnetic,\n",
    "             y_mulliken=y_mulliken,\n",
    "             y_dipole=y_dipole,\n",
    "             y_potential=y_potential,\n",
    "             m=m)\n",
    "n_types = int(x_type[~np.isnan(x_type)].max() + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_memmap = True\n",
    "load_fn = np.load if not use_memmap else partial(np.lib.format.open_memmap, mode='r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_qm9_mulliken = load_fn(f'x_qm9_mulliken{ext}.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1405126, 3, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 4, 29),\n",
       " (1405126, 9, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 3),\n",
       " (1405126, 1),\n",
       " (1405126,)]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[v.shape for v in [x_xyz,x_type,x_ext,x_atom,x_qm9_mulliken, \n",
    "                   y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential, m]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1405126, 29)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "real_atom_mask = (x_atom != -1).swapaxes(1,2).squeeze(-1)\n",
    "real_atom_mask.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_xyz_masked = x_xyz.swapaxes(1,2)[real_atom_mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(26154396, 3)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_xyz_masked.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BAD mean/std\n",
    "#x_xyz_mean, x_xyz_std = Tensor(x_xyz.mean(axis=(0,2),keepdims=True)), Tensor(x_xyz.std(axis=(0,2),keepdims=True))\n",
    "#x_xyz_mean = Tensor(x_xyz).mean(dim=(0,2),keepdim=True)\n",
    "#x_xyz_std  = Tensor(x_xyz).std(dim=(0,2), keepdim=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GOOD mean/std\n",
    "x_xyz_mean = Tensor(x_xyz_masked).mean(dim=(0),keepdim=True).unsqueeze(-1)\n",
    "x_xyz_std  = Tensor(x_xyz_masked).std(dim=(0), keepdim=True).unsqueeze(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 0.0094],\n",
       "          [-0.0328],\n",
       "          [ 0.0032]]]), tensor([[[1.8703],\n",
       "          [2.2832],\n",
       "          [1.7753]]]), torch.Size([1, 3, 1]))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_xyz_mean, x_xyz_std, x_xyz_std.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#x_xyz = (x_xyz - x_xyz_mean.numpy() ) / x_xyz_std.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_qm9_mulliken_mean = Tensor(x_qm9_mulliken.mean(axis=(0,2),keepdims=True))\n",
    "x_qm9_mulliken_std  = Tensor(x_qm9_mulliken.std( axis=(0,2),keepdims=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 1])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_xyz_std.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fastai classes (this should should be done into its own `application` but who has time?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoleculeItem(ItemBase):\n",
    "    def __init__(self,i,xyz,type,ext,atom,qm9_mulliken): \n",
    "        self.i,self.xyz,self.type,self.ext, self.atom,self.qm9_mulliken = \\\n",
    "            i,xyz,type,ext,atom,qm9_mulliken\n",
    "        self.data = [Tensor(xyz), LongTensor((type)), \n",
    "                     Tensor(ext), LongTensor((atom)),Tensor(qm9_mulliken)]\n",
    "    def __str__(self):\n",
    "        # TODO: count n_atoms correctly. \n",
    "        n_atoms = np.count_nonzero(np.sum(np.absolute(self.xyz), axis=0))+1\n",
    "        n_couplings = np.sum((self.type!=-1))\n",
    "        return f'{self.i} {n_atoms} atoms {n_couplings} couplings'\n",
    "    \n",
    "    def apply_tfms(self, tfms:Collection, **kwargs):\n",
    "        x = self.clone()\n",
    "        for t in tfms:\n",
    "            if t: x.data = t(x.data)\n",
    "        return x\n",
    "    \n",
    "    def clone(self):\n",
    "        return self.__class__(self.i,self.xyz,self.type,self.ext,self.atom,self.qm9_mulliken)\n",
    "    \n",
    "class ScalarCouplingItem(ItemBase):\n",
    "    def __init__(self,scalar,magnetic,mulliken,dipole,potential,**kwargs): \n",
    "        self.scalar,self.magnetic,self.mulliken,self.dipole,self.potential = \\\n",
    "            scalar,magnetic,mulliken,dipole,potential\n",
    "        self.data = (Tensor(scalar), Tensor(magnetic), Tensor(dipole), Tensor(potential))\n",
    "    def __str__(self):\n",
    "        res, spacer, n_couplings = '', '', 0\n",
    "        for s in self.data[0].sum(dim=0):\n",
    "            if s==0.: spacer = ' * '\n",
    "            else: \n",
    "                res += f'{spacer}{s:.4f}'\n",
    "                spacer = ' '\n",
    "                n_couplings +=1\n",
    "        return f'{n_couplings}: {res}'\n",
    "    def apply_tfms(self, tfms:Collection, **kwargs):\n",
    "        y = self.clone()\n",
    "        for t in tfms:\n",
    "            if 'label_smoothing' == t.__name__:\n",
    "                if t: y.data = t(y.data)\n",
    "        return y\n",
    "    def clone(self):\n",
    "        return self.__class__(self.scalar,self.magnetic,self.mulliken,self.dipole,self.potential)\n",
    "    def __hash__(self): return hash(str(self))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 3.588289,  0.018822,  0.085776, -0.053406], dtype=float32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_scalar.mean(axis=(0,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LMAEMaskedLoss(Module):\n",
    "    def __init__(self, contrib_w=0., magnetic_w=0., dipole_w=0., potential_w=0., \n",
    "                 types_w = [1]*n_types, return_all=False, proxy_log=torch.log, exclude_ext=False):\n",
    "        self.contrib_w,self.magnetic_w,self.dipole_w,self.potential_w = contrib_w,magnetic_w,dipole_w,potential_w\n",
    "        self.types_w = types_w\n",
    "        self.return_all = return_all\n",
    "        self.proxy_log = proxy_log\n",
    "        self.exclude_ext = exclude_ext\n",
    "    \n",
    "    def forward(self, input_outputs, t_scalar, t_magnetic, t_dipole, t_potential):    \n",
    "        type, ext, p_scalar, p_magnetic, p_dipole, p_potential = input_outputs\n",
    "        loss = 0.\n",
    "        n = 0\n",
    "        j_loss = [0] * n_types\n",
    "        for t in range(n_types):\n",
    "            mask = (type == t).squeeze(1) if not self.exclude_ext else ((type == t) & (ext == 0)).squeeze(1)\n",
    "            if mask.sum() > 0:\n",
    "                _output,_target = p_scalar.transpose(1,2)[mask], t_scalar.transpose(1,2)[mask] # scalars at the end\n",
    "                # LMAE scalar\n",
    "                s_loss = self.proxy_log((_output.sum(dim=-1) - _target.sum(dim=-1)).abs().mean()+1e-9)\n",
    "                loss += self.types_w[t] * s_loss\n",
    "                j_loss[t] += s_loss\n",
    "                # LMAE scalar contributions\n",
    "                for i_contrib in range(_output.shape[-1]):\n",
    "                    loss += self.contrib_w * ((_output[...,i_contrib] - _target[...,i_contrib])**2).mean()\n",
    "                n+=1\n",
    "        loss /= n\n",
    "        \n",
    "        if self.magnetic_w > 0:\n",
    "            mask = ~torch.isnan(t_magnetic)\n",
    "            loss += self.magnetic_w * MSELossFlat()(p_magnetic[mask], t_magnetic[mask])\n",
    "            \n",
    "        if self.dipole_w    > 0: loss += self.dipole_w    * MSELossFlat()(p_dipole,    t_dipole)\n",
    "        if self.potential_w > 0: loss += self.potential_w * MSELossFlat()(p_potential, t_potential)\n",
    "\n",
    "        return loss if not self.return_all else (loss, *j_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScalarCouplingList(ItemList):\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.loss_func = LMAEMaskedLoss\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        return ScalarCouplingItem(*o)\n",
    "\n",
    "    def reconstruct(self,t): return 0; # TODO for viz !!!! ScalarCouplingItem(t.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quaterions allow us to rotate 3d points randoming with a nice uniform distribution of 3 numbers hece we use them, however it's still to be seen if are useful here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py\n",
    "def qrot(q, v):\n",
    "    \"\"\"\n",
    "    Rotate vector(s) v about the rotation described by quaternion(s) q.\n",
    "    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,\n",
    "    where * denotes any number of dimensions.\n",
    "    Returns a tensor of shape (*, 3).\n",
    "    \"\"\"\n",
    "    assert q.shape[-1] == 4\n",
    "    assert v.shape[-1] == 3\n",
    "    assert q.shape[:-1] == v.shape[:-1]\n",
    "    \n",
    "    original_shape = list(v.shape)\n",
    "    q = q.view(-1, 4)\n",
    "    v = v.view(-1, 3)\n",
    "    \n",
    "    qvec = q[:, 1:]\n",
    "    uv = torch.cross(qvec, v, dim=1)\n",
    "    uuv = torch.cross(qvec, uv, dim=1)\n",
    "    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)\n",
    "\n",
    "def random_rotation(data):\n",
    "    x_xyz = data[0].transpose(0,1)\n",
    "    r = torch.rand(3)\n",
    "    sq1_v1,sqv1,v2_2pi,v3_2pi = torch.sqrt(1-r[:1]),torch.sqrt(r[:1]),2*math.pi*r[1:2],2*math.pi*r[2:3]\n",
    "    q = torch.cat([sq1_v1*torch.sin(v2_2pi), sq1_v1*torch.cos(v2_2pi), \n",
    "                   sqv1  *torch.sin(v3_2pi), sqv1  *torch.cos(v3_2pi)], dim=0).unsqueeze(0)\n",
    "    x_xyz = qrot(q.expand(x_xyz.shape[0],-1), x_xyz).squeeze(0).transpose(0,1)\n",
    "    return (x_xyz, *data[1:])\n",
    "\n",
    "def normalize(data):\n",
    "    sq = False\n",
    "    if data[0].ndim < 3:\n",
    "        data[0].unsqueeze_(0)\n",
    "        data[4].unsqueeze_(0)\n",
    "        sq = True\n",
    "    x_xyz      = (data[0] - x_xyz_mean)          / x_xyz_std\n",
    "    x_mulliken = (data[4] - x_qm9_mulliken_mean) / x_qm9_mulliken_std\n",
    "    if sq:\n",
    "        x_xyz.squeeze_(0)\n",
    "        x_mulliken.squeeze_(0)\n",
    "    return (x_xyz, data[1],data[2],data[3],x_mulliken)\n",
    "\n",
    "def canonize(data):\n",
    "    xyz,type,ext,atom,mulliken = data\n",
    "    mask = (atom == -1).squeeze(0)\n",
    "    i_max_atom = torch.nonzero(atom != -1).max() + 1\n",
    "    #mask_atoms = torch.ones((max_atoms,   max_atoms), dtype=torch.uint8)\n",
    "    #zeros      = torch.zeros((i_max_atom, i_max_atom), dtype=torch.uint8)\n",
    "    #mask_atoms[:zeros.shape[0],:zeros.shape[1]] = zeros\n",
    "\n",
    "    mask_atoms = torch.ones ((max_atoms, ), dtype=torch.uint8)\n",
    "    zeros      = torch.zeros((i_max_atom,), dtype=torch.uint8)\n",
    "    mask_atoms[:zeros.shape[0],] = zeros\n",
    "    #mask_atoms[(type==-1).squeeze(0),:] = 1\n",
    "    n_atoms = i_max_atom\n",
    "    xyz[:,mask], type[:,mask],ext[:,mask],atom[:,mask],mulliken[:,mask] = 0,-1,1,-1,0\n",
    "    return (xyz,type,ext,atom,mulliken, mask_atoms, n_atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "THRESHOLD=0.05\n",
    "def label_smoothing(y):\n",
    "    r = y[0] + torch.zeros_like(y[0]).uniform_(-THRESHOLD, THRESHOLD).masked_fill(y[0]==0, 0.)\n",
    "    return (r, y[1], y[2], y[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build `data` bunch etc. for fastai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = ItemList(items=(MoleculeItem(i,*v) for i,v in \n",
    "                       enumerate(zip(x_xyz,x_type,x_ext,x_atom,x_qm9_mulliken))),\n",
    "                label_cls=ScalarCouplingItem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "_, idx_valid_split = train_test_split(range(m.max()+1), test_size=0.1, random_state=13)\n",
    "idx_valid_split = np.argwhere(np.isin(m, idx_valid_split)).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "pixiedust": {
     "displayParams": {}
    }
   },
   "outputs": [],
   "source": [
    "data = data.split_by_idx(idx_valid_split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.label_from_func(\n",
    "    func=lambda o: (y_scalar[o.i], y_magnetic[o.i], y_mulliken[o.i], y_dipole[o.i], y_potential[o.i]),\n",
    "    label_cls=ScalarCouplingList)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = [normalize, canonize]\n",
    "tta_tfms = list(tfms)\n",
    "#tta_tfms.insert(0,random_rotation)\n",
    "data = data.transform((tta_tfms, tfms))#.transform_y(([label_smoothing], None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "data=data.databunch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "#data.train_ds.filter_by_func(lambda item, _: len(item.data[1].cpu().numpy()[item.data[1].cpu().numpy()==0]) == 0)\n",
    "#data.valid_ds.filter_by_func(lambda item, _: len(item.data[1].cpu().numpy()[item.data[1].cpu().numpy()==0]) == 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model\n",
    "Whole model here, self-contained (needs some cleanup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LMAEMetric(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn, val_only=True):\n",
    "        super().__init__(learn)\n",
    "        self.val_only=val_only\n",
    "        self.metric = LMAEMaskedLoss(return_all=True, exclude_ext=True)\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        if not self.val_only: self.learn.recorder.add_metric_names(['tLMAE'])\n",
    "        self.learn.recorder.add_metric_names(['👉🏻LMAE👈🏻'] + [f'{i} {types[i]}' for i in range(n_types)])\n",
    "            \n",
    "    def on_batch_end(self, train, last_output, last_target, **kwargs):\n",
    "        if self.val_only and train: return \n",
    "        preds,targs = self.preds[int(train)], self.targs[int(train)] # 0 val 1 train\n",
    "        if preds is None:\n",
    "            targs, preds = listify(last_target), listify(last_output)\n",
    "            targs,preds = [t.detach() for t in targs],[t.detach() for t in preds]\n",
    "        else:\n",
    "            for i,(o,t) in enumerate(zip(last_output, last_target)):\n",
    "                preds[i] = torch.cat([preds[i], o.detach()], dim=0)\n",
    "                targs[i] = torch.cat([targs[i], t.detach()], dim=0)\n",
    "        self.preds[int(train)], self.targs[int(train)] = preds,targs\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.targs, self.preds = [None, None], [None, None]\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        mets = []\n",
    "        if self.preds[1]: mets.append(self.metric.forward(self.preds[1], *self.targs[1])[0]) # just tLMAE\n",
    "        if self.preds[0]: mets.extend(self.metric.forward(self.preds[0], *self.targs[0]))\n",
    "        return add_metrics(last_metrics, mets) if mets else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DummyDecoder(Module):\n",
    "    def forward(self,tgt, memory, tgt_mask=None, memory_mask=None, \n",
    "                tgt_key_padding_mask=None, memory_key_padding_mask=None):\n",
    "        return memory\n",
    "\n",
    "class AtomTorchTransformer(Module):\n",
    "    def __init__(self,n_layers,n_heads,d_model,d_inner,embed_p:float=0,final_p:float=0,\n",
    "                 d_head=None,deep_decoder=False,dense_out=False, **kwargs):\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        d_head = ifnone(d_head, d_model//n_heads)        \n",
    "        \n",
    "        self.transformer = nn.Transformer(num_encoder_layers=n_layers,\n",
    "                                          nhead=n_heads,d_model=d_model,dim_feedforward=d_inner,\n",
    "                                          dropout = 0., custom_decoder=DummyDecoder())\n",
    "        \n",
    "        channels_out = d_model*n_layers if dense_out else d_model\n",
    "        channels_out_scalar = channels_out + n_types + 1\n",
    "        if deep_decoder:\n",
    "            sl = [int(channels_out_scalar/(2**d)) for d in range(int(math.ceil(np.log2(channels_out_scalar/4)-1)))]\n",
    "            self.scalar = nn.Sequential(*(list(itertools.chain.from_iterable(\n",
    "                [[nn.Conv1d(sl[i],sl[i+1],1),nn.ReLU(),nn.BatchNorm1d(sl[i+1])] for i in range(len(sl)-1)])) + \n",
    "                [nn.Conv1d(sl[-1], 4, 1)]))\n",
    "        else:\n",
    "            self.scalar = nn.Conv1d(channels_out_scalar, 4, 1)\n",
    "\n",
    "        self.magnetic  = nn.Conv1d(channels_out, 9, 1)\n",
    "        self.dipole    = nn.Linear(channels_out, 3)\n",
    "        self.potential = nn.Linear(channels_out, 1)\n",
    "        \n",
    "        self.pool = nn.AdaptiveAvgPool1d(1)\n",
    "        \n",
    "        n_atom_embedding = d_model//2\n",
    "        n_type_embedding = d_model - n_atom_embedding - 3 #- 1 - 1 -1 \n",
    "        self.type_embedding = nn.Embedding(len(types)+1,n_type_embedding)\n",
    "        self.atom_embedding = nn.Embedding(len(atoms)+1,n_atom_embedding)\n",
    "        self.drop_type, self.drop_atom = nn.Dropout(embed_p), nn.Dropout(embed_p)\n",
    "            \n",
    "    def forward(self,xyz,type,ext,atom,mulliken,mask_atoms,n_atoms):\n",
    "        bs, _, n_pts = xyz.shape        \n",
    "        t = self.drop_type(self.type_embedding((type+1).squeeze(1)))\n",
    "        a = self.drop_atom(self.atom_embedding((atom+1).squeeze(1)))\n",
    "        \n",
    "        #x = torch.cat([xyz, mulliken, ext, mask_atoms.type_as(xyz)], dim=1)\n",
    "        #x = torch.cat([xyz, mulliken, ext, (atom!=1).type_as(xyz)], dim=1)\n",
    "        x = xyz\n",
    "        x = torch.cat([x.transpose(1,2), t, a], dim=-1) * math.sqrt(self.d_model) # B,N(29),d_model\n",
    "\n",
    "        mask = mask_atoms.to(dtype=torch.bool) \n",
    "        x = x.transpose(0,1)\n",
    "        x = self.transformer(x, x, src_key_padding_mask=mask).permute(1,2,0)\n",
    "            \n",
    "        t_one_hot = torch.zeros(bs,n_types+1,n_pts,device=type.device,dtype=x.dtype).scatter_(1,type+1, 1.)\n",
    "        \n",
    "        scalar    = self.scalar(torch.cat([x, t_one_hot], dim=1))\n",
    "        #scalar    = self.scalar(x)\n",
    "\n",
    "        magnetic  = self.magnetic(x) \n",
    "        px = self.pool(x).squeeze(-1)\n",
    "        dipole    = self.dipole(px)\n",
    "        potential = self.potential(px)\n",
    "                \n",
    "        return type,ext,scalar,magnetic,dipole,potential\n",
    "    \n",
    "    def reset(self): pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This callback allows to insert multiple stateful (not averaged) metrics in one pass. Addditionally we could add metrics for train if we want to."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model instantiation: where's all your TPUs/GPUs when you need a decent hyperparam sweep?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n    def on_backward_end(self, **kwargs:Any)->None:\\n        \"Convert the gradients back to FP32 and divide them by the scale.\"\\n        if self.dynamic and grad_overflow(self.model_params) and self.loss_scale > 1:\\n            self.loss_scale /= 2\\n            self.noskip = 0\\n            #The step will be skipped since we don\\'t update the master grads so they are all None or zero\\n        else:\\n            model_g2master_g(self.model_params, self.master_params, self.flat_master)\\n            for group in self.master_params:\\n                for param in group:\\n                    if param.grad is not None: param.grad.div_(self.loss_scale)\\n            if self.clip is not None:\\n                for group in self.master_params: nn.utils.clip_grad_norm_(group, self.clip)\\n            if not self.dynamic: return\\n            self.noskip += 1\\n            if self.noskip >= self.max_noskip and self.loss_scale < self.max_scale:\\n                self.loss_scale *= 2\\n                self.noskip = 0\\n    def on_step_end(self, **kwargs:Any)->None:\\n        pass\\n        \"Update the params from master to model and zero grad.\"\\n        #Zeros the gradients of the model since the optimizer is disconnected.\\n        self.learn.model.zero_grad()\\n        #Update the params from master to model.\\n        master2model(self.model_params, self.master_params, self.flat_master)\\n'"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class ApexMixedPrecision(LearnerCallback):\n",
    "    _order = 999 #Need to run after things that could call on_backward_begin and change the loss\n",
    "    \"Callback that handles mixed-precision training.\"\n",
    "    def __init__(self, learn:Learner, loss_scale:float=None, max_noskip:int=1000, dynamic:bool=True, clip:float=None,\n",
    "                 flat_master:bool=False, max_scale:float=2**24):\n",
    "        super().__init__(learn)\n",
    "        self.flat_master,self.dynamic,self.max_noskip,self.clip,self.max_scale = flat_master,dynamic,max_noskip,clip,max_scale\n",
    "        self.loss_scale = ifnone(loss_scale, 2**16 if dynamic else 512)\n",
    "        self.not_min += ['model_params', 'master_params']\n",
    "        assert torch.backends.cudnn.enabled, \"Mixed precision training requires cudnn.\"\n",
    "        self.model,self.optimizer = amp.initialize(\n",
    "            models = learner.model, optimizers=torch.optim.Adam(learner.model.parameters(),betas=(0.9,0.99), eps=1e-4), \n",
    "            opt_level=\"O1\", keep_batchnorm_fp32=None, loss_scale=\"dynamic\")\n",
    "        #learner.create_opt(self.optimizer) \n",
    "        learner.opt = OptimWrapper(self.optimizer, wd=learner.wd, true_wd=learner.true_wd, bn_wd=learner.bn_wd)\n",
    "        learner.model = self.model\n",
    "\n",
    "    def on_train_begin(self, **kwargs:Any)->None:\n",
    "        \"Prepare the master model.\"\n",
    "        #Get a copy of the model params in FP32\n",
    "        ##self.model_params, self.master_params = get_master(self.learn.layer_groups, self.flat_master)\n",
    "        #Changes the optimizer so that the optimization step is done in FP32.\n",
    "        ##new_opt = self.learn.opt.new_with_params(self.master_params)\n",
    "        ##if self.opt is not None:\n",
    "        #    self.opt.lr,self.opt.wd = self.learn.opt.lr,self.learn.opt.wd\n",
    "        #    new_opt.load_state_dict(self.opt)\n",
    "        #self.learn.opt.opt = new_opt.opt\n",
    "        self.noskip = 0\n",
    "        \n",
    "\n",
    "    def on_loss_begin(self, last_output:Tensor, **kwargs:Any) -> Tensor:\n",
    "        \"Convert half precision output to FP32 to avoid reduction overflow.\"\n",
    "#        return {'last_output': to_float(last_output)}\n",
    "        return {'last_output': last_output}\n",
    "\n",
    "    def on_backward_begin(self, last_loss:Rank0Tensor, **kwargs:Any) -> Rank0Tensor:\n",
    "        \"Scale gradients up by `self.loss_scale` to prevent underflow.\"\n",
    "        #To avoid gradient underflow, we scale the gradients\n",
    "        #ret_loss = last_loss * self.loss_scale\n",
    "        # return {'last_loss': ret_loss}\n",
    "        with amp.scale_loss(last_loss, learner.opt.opt) as scaled_loss:\n",
    "            scaled_loss.backward()\n",
    "        return {'last_loss': scaled_loss, 'skip_bwd': True}\n",
    "\n",
    "'''\n",
    "    def on_backward_end(self, **kwargs:Any)->None:\n",
    "        \"Convert the gradients back to FP32 and divide them by the scale.\"\n",
    "        if self.dynamic and grad_overflow(self.model_params) and self.loss_scale > 1:\n",
    "            self.loss_scale /= 2\n",
    "            self.noskip = 0\n",
    "            #The step will be skipped since we don't update the master grads so they are all None or zero\n",
    "        else:\n",
    "            model_g2master_g(self.model_params, self.master_params, self.flat_master)\n",
    "            for group in self.master_params:\n",
    "                for param in group:\n",
    "                    if param.grad is not None: param.grad.div_(self.loss_scale)\n",
    "            if self.clip is not None:\n",
    "                for group in self.master_params: nn.utils.clip_grad_norm_(group, self.clip)\n",
    "            if not self.dynamic: return\n",
    "            self.noskip += 1\n",
    "            if self.noskip >= self.max_noskip and self.loss_scale < self.max_scale:\n",
    "                self.loss_scale *= 2\n",
    "                self.noskip = 0\n",
    "    def on_step_end(self, **kwargs:Any)->None:\n",
    "        pass\n",
    "        \"Update the params from master to model and zero grad.\"\n",
    "        #Zeros the gradients of the model since the optimizer is disconnected.\n",
    "        self.learn.model.zero_grad()\n",
    "        #Update the params from master to model.\n",
    "        master2model(self.model_params, self.master_params, self.flat_master)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "net, learner = None,None\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "n_layers=6\n",
    "n_heads=8\n",
    "d_model=256\n",
    "d_inner=2048\n",
    "\n",
    "deep_decoder = False\n",
    "dense_out = False\n",
    "\n",
    "net = AtomTorchTransformer(n_layers=n_layers, n_heads=n_heads,d_model=d_model,d_inner=d_inner,\n",
    "                      resid_p=0., attn_p=0., ff_p=0., embed_p=0, final_p=0.,\n",
    "                      deep_decoder=deep_decoder, dense_out=dense_out)\n",
    "\n",
    "\n",
    "learner = Learner(data,net,loss_func=LMAEMaskedLoss(),\n",
    "                  opt_func=partial(torch.optim.Adam, betas=(0.9,0.99), eps=1.0e-4)).to_fp32()\n",
    "#learner.opt_func = partial(torch.optim.Adam, betas=(0.9,0.99), eps=1.0e-4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.\n",
      "\n",
      "Defaults for this optimization level are:\n",
      "enabled                : True\n",
      "opt_level              : O1\n",
      "cast_model_type        : None\n",
      "patch_torch_functions  : True\n",
      "keep_batchnorm_fp32    : None\n",
      "master_weights         : None\n",
      "loss_scale             : dynamic\n",
      "Processing user overrides (additional kwargs that are not None)...\n",
      "After processing overrides, optimization options are:\n",
      "enabled                : True\n",
      "opt_level              : O1\n",
      "cast_model_type        : None\n",
      "patch_torch_functions  : True\n",
      "keep_batchnorm_fp32    : None\n",
      "master_weights         : None\n",
      "loss_scale             : dynamic\n"
     ]
    }
   ],
   "source": [
    "learner.callbacks.extend([\n",
    "    SaveModelCallback(learner, monitor='👉🏻LMAE👈🏻', mode='min'),\n",
    "    LMAEMetric(learner), ApexMixedPrecision(learner)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001),\n",
       " functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001),\n",
       " OptimWrapper over Adam (\n",
       " Parameter Group 0\n",
       "     amsgrad: False\n",
       "     betas: (0.9, 0.99)\n",
       "     eps: 0.0001\n",
       "     lr: 0.001\n",
       "     weight_decay: 0\n",
       " ).\n",
       " True weight decay: True)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learner.opt_func,  partial(torch.optim.Adam, betas=(0.9,0.99), eps=1.0e-4), learner.opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AtomTorchTransformer\n",
       "======================================================================\n",
       "Layer (type)         Output Shape         Param #    Trainable \n",
       "======================================================================\n",
       "Linear               [1, 2048]            526,336    True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 2048]            0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 256]             524,544    True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 2048]            526,336    True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 2048]            0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 256]             524,544    True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 2048]            526,336    True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 2048]            0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 256]             524,544    True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 2048]            526,336    True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 2048]            0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 256]             524,544    True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 2048]            526,336    True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 2048]            0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 256]             524,544    True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 2048]            526,336    True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 2048]            0          False     \n",
       "______________________________________________________________________\n",
       "Linear               [1, 256]             524,544    True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "LayerNorm            [1, 256]             512        True      \n",
       "______________________________________________________________________\n",
       "DummyDecoder         [1, 256]             0          False     \n",
       "______________________________________________________________________\n",
       "Conv1d               [4, 29]              1,064      True      \n",
       "______________________________________________________________________\n",
       "Conv1d               [9, 29]              2,313      True      \n",
       "______________________________________________________________________\n",
       "Linear               [3]                  771        True      \n",
       "______________________________________________________________________\n",
       "Linear               [1]                  257        True      \n",
       "______________________________________________________________________\n",
       "AdaptiveAvgPool1d    [256, 1]             0          False     \n",
       "______________________________________________________________________\n",
       "Embedding            [29, 125]            1,125      True      \n",
       "______________________________________________________________________\n",
       "Embedding            [29, 128]            768        True      \n",
       "______________________________________________________________________\n",
       "Dropout              [29, 125]            0          False     \n",
       "______________________________________________________________________\n",
       "Dropout              [29, 128]            0          False     \n",
       "______________________________________________________________________\n",
       "\n",
       "Total params: 6,318,234\n",
       "Total trainable params: 6,318,234\n",
       "Total non-trainable params: 0\n",
       "Optimized with 'torch.optim.adam.Adam', betas=(0.9, 0.99), eps=0.0001\n",
       "Using true weight decay as discussed in https://www.fast.ai/2018/07/02/adam-weight-decay/ \n",
       "Loss function : LMAEMaskedLoss\n",
       "======================================================================\n",
       "Callbacks functions applied \n",
       "    SaveModelCallback\n",
       "    LMAEMetric\n",
       "    ApexMixedPrecision"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learner.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sub_fname = \"bestmodel1\" # uncomment or set None to skip loading trained net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NOT loaded!  name 'sub_fname' is not defined\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(f\"Attempting to load: {sub_fname}... \", end=\"\")\n",
    "    learner.load(sub_fname, strict=False,with_opt=False)\n",
    "    print(\"Loaded\")\n",
    "except Exception as e:\n",
    "    print(\"NOT loaded! \", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# neural weight manual transplant \n",
    "if False:\n",
    "    for name,child in learner.model.named_children():\n",
    "        print(\"CHILD: \",name)\n",
    "        if not (name in ['scalar', 'magnetic', 'dipole', 'potential']):\n",
    "            print(\"FREEZING\")\n",
    "            for param in child.parameters(): param.requires_grad = False\n",
    "        else:\n",
    "            for name,param in child.named_parameters(): \n",
    "                param.requires_grad = True\n",
    "    learner.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Real loss func. Need to test different auxiliary tasks weights: `magnetic_w`, `dipole_w`, `potential`, weights of indivial `lmae`s: `types_w`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.loss_func = LMAEMaskedLoss(contrib_w = 0, magnetic_w=0, dipole_w=0, potential_w=0, \n",
    "                                  types_w = [1] * 8)#,proxy_log=lambda x: x) #\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "#learner = learner.to_parallel()#.to_fp16() #loss_scale=128, dynamic=False)\n",
    "b = 0.4\n",
    "data.batch_size = {1: int(b*(1/3)*4096//2), 2: int(b*(2/3)*4096//2), 3: int(b*4096//2)}[torch.cuda.device_count()]\n",
    "data.batch_size *= int(any([isinstance(cb, MixedPrecision) or isinstance(cb, ApexMixedPrecision) \n",
    "                            for cb in learner.callbacks]))+1 # 2x if fp16 1x otherwise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses, lrs = [], []\n",
    "wds = [None] #[1e-1,1e-2,1e-3,1e-4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SaveModelCallback\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[SaveModelCallback\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " monitor: 👉🏻LMAE👈🏻\n",
       " mode: min\n",
       " every: improvement\n",
       " name: bestmodel, LMAEMetric\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " val_only: True, ApexMixedPrecision\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " loss_scale: 65536\n",
       " max_noskip: 1000\n",
       " dynamic: True\n",
       " clip: None\n",
       " flat_master: False\n",
       " max_scale: 16777216], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " monitor: 👉🏻LMAE👈🏻\n",
       " mode: min\n",
       " every: improvement\n",
       " name: bestmodel, LMAEMetric\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[SaveModelCallback\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " monitor: 👉🏻LMAE👈🏻\n",
       " mode: min\n",
       " every: improvement\n",
       " name: bestmodel, LMAEMetric\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " val_only: True, ApexMixedPrecision\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " loss_scale: 65536\n",
       " max_noskip: 1000\n",
       " dynamic: True\n",
       " clip: None\n",
       " flat_master: False\n",
       " max_scale: 16777216], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " val_only: True, ApexMixedPrecision\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[SaveModelCallback\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " monitor: 👉🏻LMAE👈🏻\n",
       " mode: min\n",
       " every: improvement\n",
       " name: bestmodel, LMAEMetric\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " val_only: True, ApexMixedPrecision\n",
       " learn: Learner(data=DataBunch;\n",
       " \n",
       " Train: LabelList (1264961 items)\n",
       " x: ItemList\n",
       " 0 5 atoms 4 couplings,1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings\n",
       " y: ScalarCouplingList\n",
       " 4:  * 84.8076 84.8074 84.8093 84.8095,4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543\n",
       " Path: .;\n",
       " \n",
       " Valid: LabelList (140165 items)\n",
       " x: ItemList\n",
       " 14 8 atoms 6 couplings,15 8 atoms 6 couplings,16 8 atoms 7 couplings,17 8 atoms 7 couplings,18 8 atoms 7 couplings\n",
       " y: ScalarCouplingList\n",
       " 6:  * 83.5430 83.5417 83.5484 -2.3788 -2.3785 -2.3772,6:  * -2.3783 -2.3786 -2.3772 83.5418 83.5430 83.5486,7: 83.5430 -2.3783 * -11.7004 -11.6979 3.2528 13.6913 3.2521,7: 83.5417 -2.3786 -11.7004 * -11.6996 13.6924 3.2525 3.2527,7: 83.5484 -2.3772 -11.6979 -11.6996 * 3.2524 3.2524 13.6921\n",
       " Path: .;\n",
       " \n",
       " Test: None, model=AtomTorchTransformer(\n",
       "   (transformer): Transformer(\n",
       "     (encoder): TransformerEncoder(\n",
       "       (layers): ModuleList(\n",
       "         (0): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (1): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (2): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (3): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (4): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "         (5): TransformerEncoderLayer(\n",
       "           (self_attn): MultiheadAttention(\n",
       "             (out_proj): Linear(in_features=256, out_features=256, bias=True)\n",
       "           )\n",
       "           (linear1): Linear(in_features=256, out_features=2048, bias=True)\n",
       "           (dropout): Dropout(p=0.0, inplace=False)\n",
       "           (linear2): Linear(in_features=2048, out_features=256, bias=True)\n",
       "           (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "           (dropout1): Dropout(p=0.0, inplace=False)\n",
       "           (dropout2): Dropout(p=0.0, inplace=False)\n",
       "         )\n",
       "       )\n",
       "       (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "     )\n",
       "     (decoder): DummyDecoder()\n",
       "   )\n",
       "   (scalar): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (magnetic): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (dipole): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (potential): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (pool): AdaptiveAvgPool1d(output_size=1)\n",
       "   (type_embedding): Embedding(9, 125)\n",
       "   (atom_embedding): Embedding(6, 128)\n",
       "   (drop_type): Dropout(p=0, inplace=False)\n",
       "   (drop_atom): Dropout(p=0, inplace=False)\n",
       " ), opt_func=functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001), loss_func=LMAEMaskedLoss(), metrics=[], true_wd=True, bn_wd=True, wd=0.01, train_bn=True, path=PosixPath('.'), model_dir='models', callback_fns=[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False)], callbacks=[...], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " loss_scale: 65536\n",
       " max_noskip: 1000\n",
       " dynamic: True\n",
       " clip: None\n",
       " flat_master: False\n",
       " max_scale: 16777216], layer_groups=[Sequential(\n",
       "   (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (1): ParameterModule()\n",
       "   (2): ParameterModule()\n",
       "   (3): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (4): Dropout(p=0.0, inplace=False)\n",
       "   (5): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (6): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (7): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (8): Dropout(p=0.0, inplace=False)\n",
       "   (9): Dropout(p=0.0, inplace=False)\n",
       "   (10): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (11): ParameterModule()\n",
       "   (12): ParameterModule()\n",
       "   (13): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (14): Dropout(p=0.0, inplace=False)\n",
       "   (15): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (16): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (17): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (18): Dropout(p=0.0, inplace=False)\n",
       "   (19): Dropout(p=0.0, inplace=False)\n",
       "   (20): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (21): ParameterModule()\n",
       "   (22): ParameterModule()\n",
       "   (23): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (24): Dropout(p=0.0, inplace=False)\n",
       "   (25): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (26): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (27): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (28): Dropout(p=0.0, inplace=False)\n",
       "   (29): Dropout(p=0.0, inplace=False)\n",
       "   (30): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (31): ParameterModule()\n",
       "   (32): ParameterModule()\n",
       "   (33): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (34): Dropout(p=0.0, inplace=False)\n",
       "   (35): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (36): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (37): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (38): Dropout(p=0.0, inplace=False)\n",
       "   (39): Dropout(p=0.0, inplace=False)\n",
       "   (40): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (41): ParameterModule()\n",
       "   (42): ParameterModule()\n",
       "   (43): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (44): Dropout(p=0.0, inplace=False)\n",
       "   (45): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (46): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (47): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (48): Dropout(p=0.0, inplace=False)\n",
       "   (49): Dropout(p=0.0, inplace=False)\n",
       "   (50): Linear(in_features=256, out_features=256, bias=True)\n",
       "   (51): ParameterModule()\n",
       "   (52): ParameterModule()\n",
       "   (53): Linear(in_features=256, out_features=2048, bias=True)\n",
       "   (54): Dropout(p=0.0, inplace=False)\n",
       "   (55): Linear(in_features=2048, out_features=256, bias=True)\n",
       "   (56): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (57): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (58): Dropout(p=0.0, inplace=False)\n",
       "   (59): Dropout(p=0.0, inplace=False)\n",
       "   (60): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "   (61): DummyDecoder()\n",
       "   (62): Conv1d(265, 4, kernel_size=(1,), stride=(1,))\n",
       "   (63): Conv1d(256, 9, kernel_size=(1,), stride=(1,))\n",
       "   (64): Linear(in_features=256, out_features=3, bias=True)\n",
       "   (65): Linear(in_features=256, out_features=1, bias=True)\n",
       "   (66): AdaptiveAvgPool1d(output_size=1)\n",
       "   (67): Embedding(9, 125)\n",
       "   (68): Embedding(6, 128)\n",
       "   (69): Dropout(p=0, inplace=False)\n",
       "   (70): Dropout(p=0, inplace=False)\n",
       " )], add_time=True, silent=False, cb_fns_registered=False)\n",
       " loss_scale: 65536\n",
       " max_noskip: 1000\n",
       " dynamic: True\n",
       " clip: None\n",
       " flat_master: False\n",
       " max_scale: 16777216]"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learner.callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "#learner.model, learner.opt_func = learner.callbacks[-1].model, learner.callbacks[-1].opt_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "functools.partial(<class 'torch.optim.adam.Adam'>, betas=(0.9, 0.99), eps=0.0001)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "learner.opt_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rm: cannot remove 'models/bestmodel.pth': No such file or directory\r\n"
     ]
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 32768.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 16384.0\n",
      "LR Finder is complete, type {learner_name}.recorder.plot() to see the graph.\n",
      "Failed to compute the gradients, there might not be enough points.\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZQAAAEKCAYAAAA1qaOTAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvNQv5yAAAEc5JREFUeJzt3XuwXWV5x/HvT2LxgoaLQWNiDAKjE9uKZRdKbR0UCPCHBtR2oHWMSpvRylhldIrjtGi84a2o9TLNqJhaFRXHMWg1RpSxtag5EaxGxQS0JUIFDaLUC6JP/9gruj3uk7OT856zczjfz8yes9e7nrXeZ5Owf1lr7b1OqgpJkmbqHuNuQJJ092CgSJKaMFAkSU0YKJKkJgwUSVITBookqQkDRZLUhIEiSWrCQJEkNbFo3A3MpQc84AG1cuXKcbchSfPKtm3bvldVS6arW1CBsnLlSiYmJsbdhiTNK0n+e5Q6T3lJkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1MdZASXJGkuuS7Exy4ZD1Byd5f7f+C0lWTlq/IskdSV4wVz1LkoYbW6AkOQh4C3AmsAo4N8mqSWXnAbdV1THAJcCrJ62/BPj4bPcqSZreOI9QTgB2VtUNVXUncBmwZlLNGmBj9/xy4JQkAUhyFnADsH2O+pUk7cU4A2UZcOPA8q5ubGhNVd0F3A4ckeS+wN8BL52DPiVJIxhnoGTIWI1Y81Lgkqq6Y9pJknVJJpJM3HrrrfvRpiRpFIvGOPcu4CEDy8uBm6ao2ZVkEbAY2A2cCDwlyWuAQ4FfJvlpVb158iRVtQHYANDr9SYHliSpkXEGylbg2CRHAd8BzgH+YlLNJmAtcDXwFODTVVXAn+4pSPIS4I5hYSJJmjtjC5SquivJ+cBm4CDgnVW1Pcl6YKKqNgHvAN6dZCf9I5NzxtWvJGnv0v8H/8LQ6/VqYmJi3G1I0rySZFtV9aar85vykqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpoYa6AkOSPJdUl2JrlwyPqDk7y/W/+FJCu78dOSbEvyle7n4+e6d0nSbxpboCQ5CHgLcCawCjg3yapJZecBt1XVMcAlwKu78e8BT6iq3wPWAu+em64lSVMZ5xHKCcDOqrqhqu4ELgPWTKpZA2zsnl8OnJIkVXVNVd3UjW8H7pXk4DnpWpI01DgDZRlw48Dyrm5saE1V3QXcDhwxqebJwDVV9bNZ6lOSNIJFY5w7Q8ZqX2qSPJL+abDVU06SrAPWAaxYsWLfu5QkjWScRyi7gIcMLC8HbpqqJskiYDGwu1teDnwYeFpVXT/VJFW1oap6VdVbsmRJw/YlSYPGGShbgWOTHJXkd4BzgE2TajbRv+gO8BTg01VVSQ4FPga8qKo+N2cdS5KmNLZA6a6JnA9sBr4OfKCqtidZn+SJXdk7gCOS7AQuAPZ8tPh84Bjg75Nc2z2OnOOXIEkakKrJly3uvnq9Xk1MTIy7DUmaV5Jsq6redHV+U16S1ISBIklqwkCRJDVhoEiSmjBQJElNGCiSpCYMFElSEwaKJKkJA0WS1ISBIklqwkCRJDVhoEiSmjBQJElNGCiSpCYMFElSEwaKJKkJA0WS1ISBIklqwkCRJDVhoEiSmjBQJElNGCiSpCYMFElSEwaKJKkJA0WS1MRIgZLk6CQHd89PTvLcJIfObmuSpPlk1COUDwG/SHIM8A7gKOC9s9aVJGneGTVQfllVdwFnA2+oqucDS2evLUnSfDNqoPw8ybnAWuCj3dg9Z6clSdJ8NGqgPAM4CXhFVX0ryVHAv85eW5Kk+WakQKmqr1XVc6vqfUkOA+5XVRfPdPIkZyS5LsnOJBcOWX9wkvd367+QZOXAuhd149clOX2mvUiSZmbUT3ldleT+SQ4HvgxcmuQfZzJxkoOAtwBnAquAc5OsmlR2HnBbVR0DXAK8utt2FXAO8EjgDOCt3f4kSWMy6imvxVX1Q+BJwKVVdTxw6gznPgHYWVU3VNWdwGXAmkk1a4CN3fPLgVOSpBu/rKp+VlXfAnZ2+5MkjcmogbIoyVLgz/n1RfmZWgbcOLC8qxsbWtN9yux24IgRt5UkzaFRA2U9sBm4vqq2JnkYsGOGc2fIWI1YM8q2/R0k65JMJJm49dZb97FFSdKoRr0o/8Gq+v2qena3fENVPXmGc+8CHjKwvBy4aaqaJIuAxcDuEbfd0/uGqupVVW/JkiUzbFmSNJVRL8ovT/LhJLck+W6SDyVZPsO5twLHJjkqye/Qv8i+aVLNJvrffQF4CvDpqqpu/JzuU2BHAccCX5xhP5KkGRj1lNel9N/EH0z/WsUV3dh+666JnE//VNrXgQ9U1fYk65M8sSt7B3BEkp3ABcCF3bbbgQ8AXwM+ATynqn4xk34kSTOT/j/4pylKrq2q46YbO9D1er2amJgYdxuSNK8k2VZVvenqRj1C+V6SpyY5qHs8Ffj+zFqUJN2djBooz6T/keH/BW6mfz3jGbPVlCRp/hn1U17/U1VPrKolVXVkVZ1F/0uOkiQBM/uNjRc060KSNO/NJFCGfblQkrRAzSRQpv94mCRpwVi0t5VJfsTw4Ahw71npSJI0L+01UKrqfnPViCRpfpvJKS9Jkn7FQJEkNWGgSJKaMFAkSU0YKJKkJgwUSVITBookqQkDRZLUhIEiSWrCQJEkNWGgSJKaMFAkSU0YKJKkJgwUSVITBookqQkDRZLUhIEiSWrCQJEkNWGgSJKaMFAkSU0YKJKkJgwUSVITBookqYmxBEqSw5NsSbKj+3nYFHVru5odSdZ2Y/dJ8rEk30iyPcnFc9u9JGmYcR2hXAhcWVXHAld2y78hyeHARcCJwAnARQPB87qqegTwaOAxSc6cm7YlSVMZV6CsATZ2zzcCZw2pOR3YUlW7q+o2YAtwRlX9uKo+A1BVdwJfApbPQc+SpL0YV6A8sKpuBuh+HjmkZhlw48Dyrm7sV5IcCjyB/lGOJGmMFs3WjpN8CnjQkFUvHnUXQ8ZqYP+LgPcBb6qqG/bSxzpgHcCKFStGnFqStK9mLVCq6tSp1iX5bpKlVXVzkqXALUPKdgEnDywvB64aWN4A7KiqN0zTx4aull6vV3urlSTtv3Gd8toErO2erwU+MqRmM7A6yWHdxfjV3RhJXg4sBp43B71KkkYwrkC5GDgtyQ7gtG6ZJL0kbweoqt3Ay4Ct3WN9Ve1Ospz+abNVwJeSXJvkr8bxIiRJv5aqhXMWqNfr1cTExLjbkKR5Jcm2qupNV+c35SVJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJasJAkSQ1YaBIkpowUCRJTRgokqQmDBRJUhMGiiSpCQNFktSEgSJJamIsgZLk8CRbkuzofh42Rd3armZHkrVD1m9K8tXZ71iSNJ1xHaFcCFxZVccCV3bLvyHJ4cBFwInACcBFg8GT5EnAHXPTriRpOuMKlDXAxu75RuCsITWnA1uqandV3QZsAc4ASHIIcAHw8jnoVZI0gnEFygOr6maA7ueRQ2qWATcOLO/qxgBeBrwe+PFsNilJGt2i2dpxkk8BDxqy6sWj7mLIWCU5Djimqp6fZOUIfawD1gGsWLFixKklSftq1gKlqk6dal2S7yZZWlU3J1kK3DKkbBdw8sDycuAq4CTg+CTfpt//kUmuqqqTGaKqNgAbAHq9Xu37K5EkjWJcp7w2AXs+tbUW+MiQms3A6iSHdRfjVwObq+ptVfXgqloJ/AnwzanCRJI0d8YVKBcDpyXZAZzWLZOkl+TtAFW1m/61kq3dY303Jkk6AKVq4ZwF6vV6NTExMe42JGleSbKtqnrT1flNeUlSEwaKJKkJA0WS1ISBIklqwkCRJDVhoEiSmjBQJElNGCiSpCYMFElSEwaKJKkJA0WS1ISBIklqwkCRJDVhoEiSmjBQJElNGCiSpCYMFElSEwaKJKkJA0WS1ISBIklqwkCRJDVhoEiSmjBQJElNGCiSpCYMFElSE6mqcfcwZ5LcCvwAuH0/Nn8A8L22HWkvFrN/f04HsgP1NY2rr9met/X+W+1vJvvZ321n+v710KpaMl3RggoUgCQbqmrdfmw3UVW92ehJv21//5wOZAfqaxpXX7M9b+v9t9rfTPZzoL9/LcRTXleMuwGN5O7453SgvqZx9TXb87bef6v9zWQ/B+rfIWABHqHsL49QJM1XHqEceDaMuwFJ2k9z8v7lEYokqQmPUCRJTSyYQEnyziS3JPlqo/2tTbKje6wdGD8+yVeS7EzypiRpMZ+khW0O38NekeTGJHfs6z4XTKAA7wLO2NeNklyVZOWkscOBi4ATgROAi5Ic1q1+G7AOOLZ77POckjTEu5ib97ArurF9tmACpao+C+weHEtydJJPJNmW5N+TPGLE3Z0ObKmq3VV1G7AFOCPJUuD+VXV19S9O/QtwVsvXIWlhmov3sG6ez1fVzfvT46L92ehuZAPwrKrakeRE4K3A40fYbhlw48Dyrm5sWfd88rgkzYbW72EzsmADJckhwB8DHxy4zHFwt+4ZwN92Y8cA/5bkTuBbVXU2MOy6SO1lXJKamqX3sBlZsIFC/3TfD6rquMkrqupS4FLon38Enl5V3x4o2QWcPLC8HLiqG18+afymhj1L0h6z8R4244YWpKr6IfCtJH8GkL5Hjbj5ZmB1ksO6C1mrgc3deccfJfmj7tNdTwM+Mhv9S1rYZuM9bKY9LZhASfI+4Grg4Ul2JTkP+EvgvCRfBrYDa0bZV1XtBl4GbO0e67sxgGcDbwd2AtcDH2/6QiQtSHP1HpbkNUl2Affp5nnJyD36TXlJUgsL5ghFkjS7DBRJUhMGiiSpCQNFktSEgSJJasJA0YK2P3dUneF8b0+yqtG+fpHk2iRfTXJFkkOnqT80yd+0mFsaxo8Na0FLckdVHdJwf4uq6q5W+5tmrl/1nmQj8M2qesVe6lcCH62q352L/rTweIQiTZJkSZIPJdnaPR7TjZ+Q5D+TXNP9fHg3/vQkH0xyBfDJJCd3twy/PMk3krxnz+/F6cZ73fM7ut898eUkn0/ywG786G55a5L1Ix5FXU13c78khyS5MsmXut/Ns+fLbhcDR3dHNa/tal/YzfNfSV7a8D+jFiADRfptbwQuqao/BJ5M/84HAN8AHltVjwb+AXjlwDYnAWuras+dXh8NPA9YBTwMeMyQee4LfL6qHgV8Fvjrgfnf2M0/7b3gkhwEnAJs6oZ+CpxdVX8APA54fRdoFwLXV9VxVfXCJKvp/86eE4DjgOOTPHa6+aSpLOSbQ0pTORVYNXAH1/snuR+wGNiY5Fj6d2a958A2WwZuvwPwxaraBZDkWmAl8B+T5rkT+Gj3fBtwWvf8JH79e3TeC7xuij7vPbDvbfR/pwX07yT7yi4cfkn/yOWBQ7Zf3T2u6ZYPoR8wn51iPmmvDBTpt90DOKmqfjI4mOSfgM9U1dnd9YirBlb/36R9/Gzg+S8Y/v/az+vXFzGnqtmbn1TVcUkW0w+m5wBvon9/pyXA8VX18yTfBu41ZPsAr6qqf97HeaWhPOUl/bZPAufvWUiy5/bgi4HvdM+fPovzf57+qTaAc6YrrqrbgecCL0hyT/p93tKFyeOAh3alPwLuN7DpZuCZ3e/VIMmyJEc2eg1agAwULXR77qi653EB/TfnXneh+mvAs7ra1wCvSvI54KBZ7Ol5wAVJvggsBW6fboOqugb4Mv0Aeg/9/ifoH618o6v5PvC57mPGr62qT9I/pXZ1kq8Al/ObgSPtEz82LB1gktyH/umsSnIOcG5VjXRbcmmcvIYiHXiOB97cfTLrB8Azx9yPNBKPUCRJTXgNRZLUhIEiSWrCQJEkNWGgSJKaMFAkSU0YKJKkJv4fB7PKsRnI19oAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7ff0d1d03278>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for wd in wds:\n",
    "    !rm models/bestmodel.pth\n",
    "    learner.lr_find(wd=wd)\n",
    "    learner.recorder.plot(suggestion=True)\n",
    "    losses.append(learner.recorder.losses)\n",
    "    lrs.append(learner.recorder.lrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainingPhase(length=2964480)"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epochs = 1280\n",
    "b_its = len(data.train_dl)\n",
    "ph1 = (TrainingPhase(epochs*b_its)\n",
    "            .schedule_hp('wd', (1e-5,1e-1))\n",
    "            )\n",
    "gs = GeneralScheduler(learner, (ph1,))\n",
    "ph1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model bestmodel_260 not found.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='0' class='' max='1019', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      0.00% [0/1019 00:00<00:00]\n",
       "    </div>\n",
       "    \n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>👉🏻LMAE👈🏻</th>\n",
       "      <th>0 1JHC</th>\n",
       "      <th>1 2JHH</th>\n",
       "      <th>2 1JHN</th>\n",
       "      <th>3 2JHN</th>\n",
       "      <th>4 2JHC</th>\n",
       "      <th>5 3JHH</th>\n",
       "      <th>6 3JHC</th>\n",
       "      <th>7 3JHN</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>\n",
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='0' class='progress-bar-interrupted' max='2316', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      Interrupted\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8192.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2048.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1024.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 512.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 256.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 128.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 64.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 32.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 16.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.5\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.25\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.125\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.0625\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.03125\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.015625\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.0078125\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.00390625\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.001953125\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.0009765625\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.00048828125\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.000244140625\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 0.0001220703125\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.103515625e-05\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.0517578125e-05\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.52587890625e-05\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.62939453125e-06\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.814697265625e-06\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.9073486328125e-06\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.5367431640625e-07\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.76837158203125e-07\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.384185791015625e-07\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1920928955078125e-07\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.960464477539063e-08\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.9802322387695312e-08\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.4901161193847656e-08\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.450580596923828e-09\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.725290298461914e-09\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.862645149230957e-09\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.313225746154785e-10\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.656612873077393e-10\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.3283064365386963e-10\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1641532182693481e-10\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.820766091346741e-11\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.9103830456733704e-11\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.4551915228366852e-11\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.275957614183426e-12\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.637978807091713e-12\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.8189894035458565e-12\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.094947017729282e-13\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.547473508864641e-13\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.2737367544323206e-13\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1368683772161603e-13\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.684341886080802e-14\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.842170943040401e-14\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.4210854715202004e-14\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.105427357601002e-15\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.552713678800501e-15\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.7763568394002505e-15\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.881784197001252e-16\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.440892098500626e-16\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.220446049250313e-16\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1102230246251565e-16\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.551115123125783e-17\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.7755575615628914e-17\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.3877787807814457e-17\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.938893903907228e-18\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.469446951953614e-18\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.734723475976807e-18\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.673617379884035e-19\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.336808689942018e-19\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.168404344971009e-19\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0842021724855044e-19\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.421010862427522e-20\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.710505431213761e-20\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.3552527156068805e-20\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.776263578034403e-21\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.3881317890172014e-21\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.6940658945086007e-21\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.470329472543003e-22\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.235164736271502e-22\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.117582368135751e-22\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0587911840678754e-22\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.293955920339377e-23\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.6469779601696886e-23\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.3234889800848443e-23\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.617444900424222e-24\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.308722450212111e-24\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.6543612251060553e-24\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.271806125530277e-25\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.1359030627651384e-25\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.0679515313825692e-25\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0339757656912846e-25\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.169878828456423e-26\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.5849394142282115e-26\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.2924697071141057e-26\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.462348535570529e-27\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.2311742677852644e-27\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.6155871338926322e-27\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.077935669463161e-28\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.0389678347315804e-28\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.0194839173657902e-28\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0097419586828951e-28\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.048709793414476e-29\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.524354896707238e-29\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.262177448353619e-29\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.310887241768095e-30\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.1554436208840472e-30\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.5777218104420236e-30\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.888609052210118e-31\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.944304526105059e-31\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.9721522630525295e-31\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.860761315262648e-32\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.930380657631324e-32\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.465190328815662e-32\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.232595164407831e-32\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.162975822039155e-33\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.0814879110195774e-33\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.5407439555097887e-33\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.703719777548943e-34\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.851859888774472e-34\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.925929944387236e-34\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.62964972193618e-35\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.81482486096809e-35\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.407412430484045e-35\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.2037062152420224e-35\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.018531076210112e-36\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.009265538105056e-36\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.504632769052528e-36\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.52316384526264e-37\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.76158192263132e-37\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.88079096131566e-37\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.4039548065783e-38\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.70197740328915e-38\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.350988701644575e-38\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1754943508222875e-38\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.877471754111438e-39\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.938735877055719e-39\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.4693679385278594e-39\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.346839692639297e-40\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.6734198463196485e-40\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.8367099231598242e-40\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 9.183549615799121e-41\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.591774807899561e-41\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.2958874039497803e-41\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1479437019748901e-41\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.739718509874451e-42\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.8698592549372254e-42\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.4349296274686127e-42\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.174648137343064e-43\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.587324068671532e-43\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.793662034335766e-43\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.96831017167883e-44\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.484155085839415e-44\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.2420775429197073e-44\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.1210387714598537e-44\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.605193857299268e-45\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.802596928649634e-45\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.401298464324817e-45\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 7.006492321624085e-46\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.503246160812043e-46\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.7516230804060213e-46\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.758115402030107e-47\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.3790577010150533e-47\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.1895288505075267e-47\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0947644252537633e-47\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.473822126268817e-48\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.7369110631344083e-48\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.3684555315672042e-48\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.842277657836021e-49\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.4211388289180104e-49\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.7105694144590052e-49\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.552847072295026e-50\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.276423536147513e-50\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.1382117680737565e-50\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.0691058840368783e-50\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.345529420184391e-51\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.6727647100921956e-51\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.3363823550460978e-51\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 6.681911775230489e-52\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 3.3409558876152446e-52\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.6704779438076223e-52\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 8.352389719038111e-53\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4.176194859519056e-53\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.088097429759528e-53\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1.044048714879764e-53\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 5.22024357439882e-54\n",
      "Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 2.61012178719941e-54\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-56-ae2fef03d221>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mlearner\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit_one_cycle\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepochs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1e-4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpct_start\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m261\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;31m#, moms=(0.95,0.85), pct_start=0.3)#,callbacks=gs)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/train.py\u001b[0m in \u001b[0;36mfit_one_cycle\u001b[0;34m(learn, cyc_len, max_lr, moms, div_factor, pct_start, final_div, wd, callbacks, tot_epochs, start_epoch)\u001b[0m\n\u001b[1;32m     20\u001b[0m     callbacks.append(OneCycleScheduler(learn, max_lr, moms=moms, div_factor=div_factor, pct_start=pct_start,\n\u001b[1;32m     21\u001b[0m                                        final_div=final_div, tot_epochs=tot_epochs, start_epoch=start_epoch))\n\u001b[0;32m---> 22\u001b[0;31m     \u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcyc_len\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_lr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwd\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mwd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mlr_find\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlearn\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mLearner\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_lr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mFloats\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1e-7\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mend_lr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mFloats\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_it\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstop_div\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mbool\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwd\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/basic_train.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, epochs, lr, wd, callbacks)\u001b[0m\n\u001b[1;32m    205\u001b[0m         \u001b[0mcallbacks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mcb\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mcb\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallback_fns\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mlistify\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdefaults\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextra_callback_fns\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mlistify\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    206\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcb_fns_registered\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 207\u001b[0;31m         \u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepochs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmetrics\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmetrics\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    208\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    209\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mcreate_opt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mFloats\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwd\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mFloats\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m->\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/basic_train.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(epochs, learn, callbacks, metrics)\u001b[0m\n\u001b[1;32m    104\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0myb\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mprogress_bar\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_dl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparent\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpbar\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    105\u001b[0m                 \u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0myb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_batch_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0myb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 106\u001b[0;31m                 \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0myb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_func\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlearn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    107\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_batch_end\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    108\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/basic_train.py\u001b[0m in \u001b[0;36mloss_batch\u001b[0;34m(model, xb, yb, loss_func, opt, cb_handler)\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mopt\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 35\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mskip_bwd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_backward_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     36\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mskip_bwd\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mamp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopt\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mscaled_loss\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/callback.py\u001b[0m in \u001b[0;36mon_backward_begin\u001b[0;34m(self, loss)\u001b[0m\n\u001b[1;32m    290\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msmoothener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_value\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    291\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'last_loss'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'smooth_loss'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msmoothener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msmooth\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 292\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'backward_begin'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcall_mets\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    293\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'last_loss'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'skip_bwd'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    294\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/callback.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, cb_name, call_mets, **kwargs)\u001b[0m\n\u001b[1;32m    249\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mcall_mets\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    250\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmet\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmetrics\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_and_update\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmet\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcb_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 251\u001b[0;31m         \u001b[0;32mfor\u001b[0m \u001b[0mcb\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_and_update\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcb_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    252\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    253\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mset_dl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdl\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mDataLoader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/fastai/callback.py\u001b[0m in \u001b[0;36m_call_and_update\u001b[0;34m(self, cb, cb_name, **kwargs)\u001b[0m\n\u001b[1;32m    239\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_and_update\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcb_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m->\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    240\u001b[0m         \u001b[0;34m\"Call `cb_name` on `cb` and update the inner state.\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 241\u001b[0;31m         \u001b[0mnew\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mifnone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf'on_{cb_name}'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    242\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnew\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    243\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mk\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-40-cfcb8ca8bd57>\u001b[0m in \u001b[0;36mon_backward_begin\u001b[0;34m(self, last_loss, **kwargs)\u001b[0m\n\u001b[1;32m     40\u001b[0m         \u001b[0;31m# return {'last_loss': ret_loss}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     41\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mamp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlast_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlearner\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopt\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mscaled_loss\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 42\u001b[0;31m             \u001b[0mscaled_loss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     43\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m'last_loss'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mscaled_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'skip_bwd'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     44\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    116\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    117\u001b[0m         \"\"\"\n\u001b[0;32m--> 118\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    119\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    120\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     91\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     92\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 93\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "learner.fit_one_cycle(epochs,1e-4, pct_start=0.3, start_epoch=0)#, moms=(0.95,0.85), pct_start=0.3)#,callbacks=gs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.reco1rder.plot_losses()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.recorder.plot_lr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.recorder.plot_metrics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lookahead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from itertools import chain\n",
    "import warnings\n",
    "\n",
    "class Lookahead(torch.optim.Optimizer):\n",
    "    def __init__(self, optimizer, k=5, alpha=0.5):\n",
    "        self.optimizer = optimizer\n",
    "        self.k = k\n",
    "        self.alpha = alpha\n",
    "        self.param_groups = self.optimizer.param_groups\n",
    "        self.state = defaultdict(dict)\n",
    "        self.fast_state = self.optimizer.state\n",
    "        for group in self.param_groups:\n",
    "            group[\"counter\"] = 0\n",
    "    \n",
    "    def update(self, group):\n",
    "        for fast in group[\"params\"]:\n",
    "            param_state = self.state[fast]\n",
    "            if \"slow_param\" not in param_state:\n",
    "                param_state[\"slow_param\"] = torch.zeros_like(fast.data)\n",
    "                param_state[\"slow_param\"].copy_(fast.data)\n",
    "            slow = param_state[\"slow_param\"]\n",
    "            slow += (fast.data - slow) * self.alpha\n",
    "            fast.data.copy_(slow)\n",
    "    \n",
    "    def update_lookahead(self):\n",
    "        for group in self.param_groups:\n",
    "            self.update(group)\n",
    "\n",
    "    def step(self, closure=None):\n",
    "        loss = self.optimizer.step(closure)\n",
    "        for group in self.param_groups:\n",
    "            if group[\"counter\"] == 0:\n",
    "                self.update(group)\n",
    "            group[\"counter\"] += 1\n",
    "            if group[\"counter\"] >= self.k:\n",
    "                group[\"counter\"] = 0\n",
    "        return loss\n",
    "\n",
    "    def state_dict(self):\n",
    "        fast_state_dict = self.optimizer.state_dict()\n",
    "        slow_state = {\n",
    "            (id(k) if isinstance(k, torch.Tensor) else k): v\n",
    "            for k, v in self.state.items()\n",
    "        }\n",
    "        fast_state = fast_state_dict[\"state\"]\n",
    "        param_groups = fast_state_dict[\"param_groups\"]\n",
    "        return {\n",
    "            \"fast_state\": fast_state,\n",
    "            \"slow_state\": slow_state,\n",
    "            \"param_groups\": param_groups,\n",
    "        }\n",
    "\n",
    "    def load_state_dict(self, state_dict):\n",
    "        slow_state_dict = {\n",
    "            \"state\": state_dict[\"slow_state\"],\n",
    "            \"param_groups\": state_dict[\"param_groups\"],\n",
    "        }\n",
    "        fast_state_dict = {\n",
    "            \"state\": state_dict[\"fast_state\"],\n",
    "            \"param_groups\": state_dict[\"param_groups\"],\n",
    "        }\n",
    "        super(Lookahead, self).load_state_dict(slow_state_dict)\n",
    "        self.optimizer.load_state_dict(fast_state_dict)\n",
    "        self.fast_state = self.optimizer.state\n",
    "\n",
    "    def add_param_group(self, param_group):\n",
    "        param_group[\"counter\"] = 0\n",
    "        self.optimizer.add_param_group(param_group)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional Fine tune regular fit "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.callbacks.append(\n",
    "    ReduceLROnPlateauCallback(learner, monitor='train_loss', mode='min', factor=0.2, patience=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    base_opt\n",
    "except:\n",
    "    base_opt = learner.opt.opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lookahead = Lookahead(base_opt, k=5, alpha=0.5) # Initialize Lookahead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.opt = OptimWrapper(lookahead)\n",
    "learner.opt_func = lookahead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10, 4e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(50, 3e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(50, 1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(50, 8e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10, 6e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10, 5e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10, 3e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10, 8e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.recorder.plot_lr(show_moms=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.recorder.plot_losses()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.recorder.plot_metrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10,1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.to_fp32()\n",
    "val = learner.validate()[1]\n",
    "print(val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    sub_fname = f'loss{learner.recorder.losses[-1]:.04f}val{val:.04f}'\n",
    "except:\n",
    "    sub_fname = f'val{val:.04f}'\n",
    "learner.save(sub_fname)\n",
    "print(sub_fname)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference\n",
    "\n",
    "Make sure `tranforms` are activated to test set otherwise TTA > 1 will be as TTA =1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fname = Path('test.npz')\n",
    "try:\n",
    "    npzfile  = np.load(fname_ext(test_fname, ext))\n",
    "    xt_xyz   = npzfile['x_xyz']\n",
    "    xt_type  = npzfile['x_type']\n",
    "    xt_ext   = npzfile['x_ext']\n",
    "    xt_atom  = npzfile['x_atom']\n",
    "    mt = npzfile['m']\n",
    "    xt_ids = npzfile['x_ids']\n",
    "except:\n",
    "    xt_xyz,xt_type,xt_ext,xt_atom,mt,xt_ids = \\\n",
    "        preprocess(test_fname.with_suffix('.csv'), type_index=types,ext=ext)\n",
    "    np.savez(fname_ext('_'+test_fname, ext), \n",
    "             x_xyz  = xt_xyz,\n",
    "             x_type = xt_type,\n",
    "             x_ext  = xt_ext,\n",
    "             x_atom = xt_atom,\n",
    "             m=mt,\n",
    "             x_ids=xt_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xt_qm9_mulliken = load_fn(f'xt_qm9_mulliken{ext}.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[v.shape for v in [xt_xyz,xt_type,xt_ext,xt_atom, xt_qm9_mulliken,xt_ids, mt]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.data.add_test(ItemList(items=(MoleculeItem(i,*v) for i,v in \n",
    "                              enumerate(zip(xt_xyz,xt_type,xt_ext,xt_atom,xt_qm9_mulliken)))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TTA_N = 1\n",
    "learner.data.test_ds.tfms = tta_tfms if TTA_N > 1 else tfms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#learner.model = learner.model.module\n",
    "#data.batch_size = 4096*2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#learner.model = nn.DataParallel(learner.model)\n",
    "old_bs = data.batch_size\n",
    "data.batch_size *= 2\n",
    "\n",
    "sub = defaultdict(int)\n",
    "xt_ids_not_extended = (xt_ids!=0) & (xt_ids<=7163688) # TODO\n",
    "ids = xt_ids[xt_ids_not_extended]\n",
    "\n",
    "mb = master_bar(range(TTA_N))\n",
    "for tta in mb:\n",
    "    test_preds = np.zeros((0, 29), dtype=np.float32)\n",
    "\n",
    "    for batch_idx, batch in progress_bar(\n",
    "        enumerate(learner.dl(DatasetType.Test)), total=len(learner.dl(DatasetType.Test)), parent=mb):\n",
    "        _, _, preds_,_,_,_ = learner.pred_batch(ds_type=DatasetType.Test, batch=batch)\n",
    "        preds_ = preds_.sum(dim=1)\n",
    "        test_preds = np.concatenate([test_preds, preds_.data.cpu().numpy()], axis = 0)\n",
    "\n",
    "    preds = test_preds[xt_ids_not_extended]\n",
    "    for k in range(len(ids)):\n",
    "        sub[int(ids[k])] += preds[k]\n",
    "    \n",
    "for k in range(len(ids)):\n",
    "    sub[int(ids[k])] = sub[int(ids[k])]/TTA_N\n",
    "\n",
    "learner.model = learner.model.module\n",
    "data.batch_size = old_bs\n",
    "\n",
    "sub_df = pd.DataFrame(sub.items(), columns=['id', 'scalar_coupling_constant'])\n",
    "sub_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit to Kaggle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csv_fname = f'{sub_fname}_tta{TTA_N}.csv'\n",
    "sub_df.to_csv(csv_fname, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp = 'champs-scalar-coupling'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kaggle competitions submit -c {comp} -f {csv_fname} -m 'QM9 tta {TTA_N} {ext}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "time.sleep(60)\n",
    "!kaggle competitions submissions -c {comp} -v > submissions-{comp}.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "submissions = pd.read_csv(f'submissions-{comp}.csv')\n",
    "submissions.iloc[0].publicScore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}