IRON_wangj.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "###############################  part-1 : import lib ###########################\n",
    "import numpy as np\n",
    "from skimage import io\n",
    "from skimage import transform\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import torch\n",
    "from   torch import optim\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision.transforms import transforms\n",
    "from torchvision.utils import make_grid\n",
    "import csv\n",
    "import re\n",
    "from torch.autograd import Variable\n",
    "from fastai.basics import *\n",
    "\n",
    "acc_threshold = 1/22\n",
    "G_inChans  = 1\n",
    "G_outChans = 3\n",
    "\n",
    "#export\n",
    "##############################  part-1 : dataset  ###########################\n",
    "def get_file_path(root_path,file_list,dir_list):\n",
    "    dir_or_files = os.listdir(root_path)\n",
    "    for dir_file in dir_or_files:\n",
    "        dir_file_path = os.path.join(root_path,dir_file)\n",
    "        if os.path.isdir(dir_file_path):\n",
    "            dir_list.append(dir_file_path)\n",
    "            get_file_path(dir_file_path,file_list,dir_list)\n",
    "        else:\n",
    "            file_list.append(dir_file_path)\n",
    "            \n",
    "class jax_Dataset(Dataset):\n",
    "    def __init__(self, root_dir, transform=None):\n",
    "        super(jax_Dataset,self).__init__()\n",
    "        self.root_dir  = root_dir\n",
    "        self.transform = transform\n",
    "        self.size      = 0\n",
    "        self.csv_list      = []\n",
    "        self.npyfile_list  = []\n",
    "        self.label_list    = []\n",
    "        self.all_file_list = [] \n",
    "        self.dir_list      = []\n",
    "        get_file_path(self.root_dir,self.all_file_list,self.dir_list)\n",
    "        for f in self.dir_list: \n",
    "            self.csv_list.append(f+'/label.csv')    \n",
    "        for csv_path in self.csv_list:\n",
    "            with open(csv_path, 'r') as csvfile:\n",
    "                header_row = next(csvfile)\n",
    "                remain_rows = csv.reader(csvfile)\n",
    "                for row in remain_rows:\n",
    "                    d_path = row[0]\n",
    "                    d_labels = row[1:G_outChans+1]\n",
    "                    d_labels_float = [float(item) for item in d_labels] # float list\n",
    "                    d_labels_torch = torch.from_numpy(np.array(d_labels_float))\n",
    "                    self.npyfile_list.append(root_dir+'/'+d_path)\n",
    "                    self.label_list.append(d_labels_torch)\n",
    "                    self.size += 1             \n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "    def __getitem__(self, idx):\n",
    "        if idx > self.size:\n",
    "            print(\"dataset out of index\")\n",
    "            return None\n",
    "        mat_path = self.npyfile_list[idx]\n",
    "        if not os.path.isfile(mat_path):\n",
    "            print(mat_path + 'does not exist!')\n",
    "            return None\n",
    "        tensor_6d = np.load(mat_path)\n",
    "        label_6d  = self.label_list[idx]\n",
    "        if self.transform:\n",
    "            tensor_6d = self.transform(tensor_6d)\n",
    "        return tensor_6d,label_6d\n",
    "\n",
    "def get_dls(train_ds, valid_ds, bs, **kwargs):\n",
    "    return (DataLoader(train_ds, batch_size=bs, shuffle=True, **kwargs),\n",
    "            DataLoader(valid_ds, batch_size=bs*2, **kwargs))\n",
    "\n",
    "class DataBunch():\n",
    "    def __init__(self, train_dl, valid_dl, c=None):\n",
    "        self.train_dl,self.valid_dl,self.c = train_dl,valid_dl,c\n",
    "    @property\n",
    "    def train_ds(self): return self.train_dl.dataset\n",
    "    @property\n",
    "    def valid_ds(self): return self.valid_dl.dataset\n",
    "    \n",
    "#export\n",
    "##############################  part-2 : model  ###########################\n",
    "\n",
    "def accuracy(out,yb): \n",
    "    tmp = (torch.abs(out-yb) <= acc_threshold).float()\n",
    "    return tmp.mean()\n",
    "def errsmean(out,yb): \n",
    "    tmp = torch.abs(out-yb)\n",
    "    indx = (tmp > acc_threshold).float()\n",
    "    largpart = indx*tmp\n",
    "    return largpart.sum() / indx.sum()\n",
    "\n",
    "\n",
    "def flatten(x):      return x.view(x.shape[0], -1)\n",
    "def linearMap(x):\n",
    "    xmax = torch.max(x)\n",
    "    xmin = torch.min(x)\n",
    "    if abs(xmax-xmin):\n",
    "        xret = (x-xmin)/(xmax-xmin)\n",
    "    else:\n",
    "        xret = torch.zeros(x.shape)\n",
    "    return xret\n",
    "\n",
    "class Lambda(nn.Module):\n",
    "    def __init__(self, func):\n",
    "        super().__init__()\n",
    "        self.func = func\n",
    "    def forward(self, x): return self.func(x)\n",
    "    \n",
    "def get_model_layers_3d(inChans = G_inChans,outChans = G_outChans):\n",
    "    layers = [  nn.Sequential( nn.Conv3d( in_channels=inChans,out_channels=64,kernel_size=3,stride=1,\n",
    "                                          padding=0,groups=1,bias=True,padding_mode='zeros'  )  , #添加 maxpooling\n",
    "                                          nn.BatchNorm3d(64),nn.ReLU()),  # 7*7\n",
    "                nn.Sequential( nn.Conv3d( in_channels=64,out_channels=128,kernel_size=3,stride=1,\n",
    "                                          padding=0,groups=1,bias=True,padding_mode='zeros'  )  , \n",
    "                                          nn.BatchNorm3d(128), nn.ReLU()), # 5*5 \n",
    "                nn.Sequential( nn.Conv3d( in_channels=128,out_channels=256,kernel_size=3,stride=1,\n",
    "                                          padding=0,groups=1,bias=True,padding_mode='zeros'  )  , \n",
    "                                          nn.BatchNorm3d(256),nn.ReLU()), # 3*3           \n",
    "                nn.Sequential( nn.Conv3d( in_channels=256,out_channels=512,kernel_size=3,stride=1,\n",
    "                                          padding=0,groups=1,bias=True,padding_mode='zeros'  )  ,\n",
    "                                          nn.BatchNorm3d(512), nn.ReLU() ), # 1*1\n",
    "                Lambda(flatten)   ,\n",
    "                nn.Linear(512,256),\n",
    "                nn.ReLU()         ,\n",
    "                nn.Linear(256,64),\n",
    "                nn.ReLU()         ,\n",
    "                nn.Linear(64,16),\n",
    "                nn.ReLU()         ,\n",
    "                nn.Linear(16,outChans)\n",
    "              ]\n",
    "    return layers\n",
    "\n",
    "##############################  part-2 : model-init ###########################\n",
    "def model_init(model):\n",
    "    for l in model:\n",
    "        if isinstance(l, nn.Sequential):\n",
    "            nn.init.kaiming_normal_(l[0].weight)\n",
    "            \n",
    "#export\n",
    "from typing import *\n",
    "def listify(o):\n",
    "    if o is None: return []\n",
    "    if isinstance(o, list): return o\n",
    "    if isinstance(o, str): return [o]\n",
    "    if isinstance(o, Iterable): return list(o)\n",
    "    return [o]\n",
    "#export\n",
    "##########################  Optimizer ########################\n",
    "def compose(x, funcs, *args, order_key='_order', **kwargs):\n",
    "    key = lambda o: getattr(o, order_key, 0)\n",
    "    for f in sorted(listify(funcs), key=key): x = f(x, **kwargs)\n",
    "    return x\n",
    "  \n",
    "def sgd_step(p, lr, **kwargs):\n",
    "    p.data.add_(-lr, p.grad.data)\n",
    "    return p\n",
    "\n",
    "#weight decay#\n",
    "def weight_decay(p, lr, wd, **kwargs):\n",
    "    p.data.mul_(1 - lr*wd)\n",
    "    return p\n",
    "weight_decay._defaults = dict(wd=0.)\n",
    "\n",
    "def l2_reg(p, lr, wd, **kwargs):\n",
    "    p.grad.data.add_(wd, p.data)\n",
    "    return p\n",
    "l2_reg._defaults = dict(wd=0.)\n",
    "\n",
    "def maybe_update(os, dest, f):\n",
    "    for o in os:\n",
    "        for k,v in f(o).items():\n",
    "            if k not in dest: dest[k] = v\n",
    "\n",
    "def get_defaults(d): return getattr(d,'_defaults',{})\n",
    "\n",
    "class Optimizer():\n",
    "    def __init__(self, params, steppers, **defaults):\n",
    "        self.steppers = listify(steppers)\n",
    "        maybe_update(self.steppers, defaults, get_defaults)\n",
    "        self.param_groups = list(params)\n",
    "        \n",
    "        if not isinstance(self.param_groups[0], list): self.param_groups = [self.param_groups]\n",
    "        self.hypers = [{**defaults} for p in self.param_groups]\n",
    "\n",
    "    def grad_params(self):\n",
    "        return [(p,hyper) for pg,hyper in zip(self.param_groups,self.hypers)\n",
    "            for p in pg if p.grad is not None]\n",
    "\n",
    "    def zero_grad(self):\n",
    "        for p,hyper in self.grad_params():\n",
    "            p.grad.detach_()\n",
    "            p.grad.zero_()\n",
    "\n",
    "    def step(self):\n",
    "        for p,hyper in self.grad_params(): compose(p, self.steppers, **hyper)\n",
    "            \n",
    "def test(a,b,cmp,cname=None):\n",
    "    if cname is None: cname=cmp.__name__\n",
    "    assert cmp(a,b),f\"{cname}:\\n{a}\\n{b}\"\n",
    "\n",
    "def test_eq(a,b): test(a,b,operator.eq,'==')\n",
    "    \n",
    "sgd_opt = partial(Optimizer, steppers=[weight_decay, sgd_step])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "######################  Merge  Runner && Learner to a new :Learner ############\n",
    "def param_getter(m): return m.parameters()\n",
    "\n",
    "class Learner():\n",
    "    def __init__(self, model, data, loss_func, cb_funcs, opt_func, lr, splitter=param_getter,cbs=None ):\n",
    "        self.model,self.data,self.loss_func = model,data,loss_func\n",
    "        self.opt_func,self.lr,self.splitter = opt_func,lr,splitter\n",
    "        self.in_train,self.logger,self.opt  = False,print,None\n",
    "        self.cbs = []\n",
    "        self.add_cb(TrainEvalCallback())\n",
    "        self.add_cbs(cbs)\n",
    "        self.add_cbs(cbf() for cbf in listify(cb_funcs))\n",
    "\n",
    "    def add_cbs(self, cbs):\n",
    "        for cb in listify(cbs): self.add_cb(cb)\n",
    "            \n",
    "    def add_cb(self, cb):\n",
    "        cb.set_runner(self) \n",
    "        setattr(self, cb.name, cb)\n",
    "        self.cbs.append(cb)     \n",
    "\n",
    "    def remove_cbs(self, cbs):\n",
    "        for cb in listify(cbs): self.cbs.remove(cb)\n",
    "            \n",
    "    def one_batch(self, i, xb, yb):\n",
    "        try:\n",
    "            self.iter = i\n",
    "            xb = xb.unsqueeze(1)          \n",
    "            xb = xb.float()\n",
    "            yb = yb.float()\n",
    "            self.xb,self.yb = xb,yb;                        self('begin_batch') \n",
    "            self.pred = self.model(self.xb);                self('after_pred')\n",
    "            self.loss = self.loss_func(self.pred, self.yb);\n",
    "            self('after_loss') \n",
    "            if not self.in_train: return\n",
    "            self.loss.backward();                           self('after_backward')\n",
    "            self.opt.step();                                self('after_step')\n",
    "            self.opt.zero_grad()\n",
    "        except CancelBatchException:                        self('after_cancel_batch')\n",
    "        finally:                                            self('after_batch')\n",
    "\n",
    "    def all_batches(self):\n",
    "        self.iters = len(self.dl)\n",
    "        try:\n",
    "            for i,(xb,yb) in enumerate(self.dl): self.one_batch(i, xb, yb)\n",
    "        except CancelEpochException: self('after_cancel_epoch')\n",
    "\n",
    "    def do_begin_fit(self, epochs):\n",
    "        self.epochs,self.loss = epochs,tensor(0.)\n",
    "        self('begin_fit')\n",
    "\n",
    "    def do_begin_epoch(self, epoch):\n",
    "        self.epoch,self.dl = epoch,self.data.train_dl\n",
    "        return self('begin_epoch')\n",
    "\n",
    "    def fit(self, epochs, cbs=None, reset_opt=False):\n",
    "        self.add_cbs(cbs)\n",
    "        if reset_opt or not self.opt: self.opt = self.opt_func(self.splitter(self.model), lr=self.lr)\n",
    "        try:\n",
    "            self.do_begin_fit(epochs)\n",
    "            for epoch in range(epochs):\n",
    "                if not self.do_begin_epoch(epoch): self.all_batches()\n",
    "\n",
    "                with torch.no_grad(): \n",
    "                    self.dl = self.data.valid_dl\n",
    "                    if not self('begin_validate'): self.all_batches()\n",
    "                self('after_epoch')\n",
    "            \n",
    "        except CancelTrainException: self('after_cancel_train')\n",
    "        finally:\n",
    "            self('after_fit')\n",
    "            self.remove_cbs(cbs)\n",
    "\n",
    "    ALL_CBS = {'begin_batch', 'after_pred', 'after_loss', 'after_backward', 'after_step',\n",
    "        'after_cancel_batch', 'after_batch', 'after_cancel_epoch', 'begin_fit',\n",
    "        'begin_epoch', 'begin_validate', 'after_epoch',\n",
    "        'after_cancel_train', 'after_fit'}\n",
    "    \n",
    "    def __call__(self, cb_name):                         \n",
    "        res = False\n",
    "        assert cb_name in self.ALL_CBS\n",
    "        for cb in sorted(self.cbs, key=lambda x: x._order): res = cb(cb_name) and res\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import re\n",
    "_camel_re1 = re.compile('(.)([A-Z][a-z]+)')\n",
    "_camel_re2 = re.compile('([a-z0-9])([A-Z])')\n",
    "def camel2snake(name):\n",
    "    s1 = re.sub(_camel_re1, r'\\1_\\2', name)\n",
    "    return re.sub(_camel_re2, r'\\1_\\2', s1).lower()\n",
    "\n",
    "class Callback():\n",
    "    _order=0  \n",
    "    def set_runner(self, run): self.run=run\n",
    "    def __getattr__(self, k):\n",
    "        try:\n",
    "            return getattr(self.run, k)\n",
    "        except KeyError:\n",
    "            raise AttributeError(r\"Jax said Clbk: object has no attribute '%s'\" % k)\n",
    "    @property\n",
    "    def name(self): \n",
    "        name = re.sub(r'Callback$', '', self.__class__.__name__)\n",
    "        return camel2snake(name or 'callback') \n",
    "\n",
    "    def __call__(self, cb_name):\n",
    "        f = getattr(self, cb_name, None)\n",
    "        if f and f(): return True\n",
    "        return False\n",
    "    \n",
    "class TrainEvalCallback(Callback):\n",
    "    _order = 3\n",
    "    def begin_fit(self):\n",
    "        self.run.n_epochs=0.\n",
    "        self.run.n_iter=0\n",
    "    \n",
    "    def after_batch(self):\n",
    "        if not self.in_train: return\n",
    "        self.run.n_epochs += 1./self.iters\n",
    "        self.run.n_iter   += 1\n",
    "        \n",
    "    def begin_epoch(self):\n",
    "        self.run.n_epochs=self.epoch\n",
    "        self.model.train() \n",
    "        self.run.in_train=True\n",
    "\n",
    "    def begin_validate(self):\n",
    "        self.model.eval() \n",
    "        self.run.in_train=False\n",
    "\n",
    "class CancelTrainException(Exception): pass                                      \n",
    "class CancelEpochException(Exception): pass\n",
    "class CancelBatchException(Exception): pass\n",
    "\n",
    "import time\n",
    "from fastprogress.fastprogress import master_bar, progress_bar\n",
    "from fastprogress.fastprogress import format_time\n",
    "\n",
    "class AvgStats():\n",
    "    def __init__(self, metrics, in_train): self.metrics,self.in_train = listify(metrics),in_train\n",
    "    def reset(self):\n",
    "        self.tot_loss,self.count = 0.,0\n",
    "        self.tot_mets = [0.] * len(self.metrics)\n",
    "    @property\n",
    "    def all_stats(self): return [self.tot_loss.item()] + self.tot_mets\n",
    "    @property\n",
    "    def avg_stats(self): return [o/self.count for o in self.all_stats] \n",
    "                                  \n",
    "    def __repr__(self):\n",
    "        if not self.count: return \"\"\n",
    "        return f\"{'train' if self.in_train else 'valid'}: {self.avg_stats}\"\n",
    "\n",
    "    def accumulate(self, run):\n",
    "        bn = run.xb.shape[0] \n",
    "        self.tot_loss += run.loss * bn \n",
    "        self.count += bn\n",
    "        for i,m in enumerate(self.metrics): \n",
    "            self.tot_mets[i] += m(run.pred, run.yb) * bn \n",
    "\n",
    "class AvgStatsCallback(Callback):\n",
    "    def __init__(self, metrics):\n",
    "        self.train_stats,self.valid_stats = AvgStats(metrics,True),AvgStats(metrics,False)\n",
    "    \n",
    "    def begin_fit(self):\n",
    "        met_names = ['loss'] + [m.__name__ for m in self.train_stats.metrics]\n",
    "        names = ['epoch'] + [f'train_{n}' for n in met_names] + [\n",
    "            f'valid_{n}' for n in met_names] + ['time'] + ['lr']\n",
    "        self.logger(names) \n",
    "    \n",
    "    def begin_epoch(self):\n",
    "        self.train_stats.reset()\n",
    "        self.valid_stats.reset()\n",
    "        self.start_time = time.time()\n",
    "        \n",
    "    def after_loss(self):\n",
    "        stats = self.train_stats if self.in_train else self.valid_stats\n",
    "        with torch.no_grad(): stats.accumulate(self.run) \n",
    "    \n",
    "    def after_epoch(self):\n",
    "        stats = [str(self.epoch)] \n",
    "        for o in [self.train_stats, self.valid_stats]:\n",
    "            stats += [f'{v:.6f}' for v in o.avg_stats] \n",
    "        stats += [format_time(time.time() - self.start_time)] \n",
    "        stats += [self.lr]\n",
    "        self.logger(stats)\n",
    "        \n",
    "class ProgressCallback(Callback):\n",
    "    _order=-1\n",
    "    def begin_fit(self):\n",
    "        self.mbar = master_bar(range(self.epochs))\n",
    "        self.mbar.on_iter_begin()\n",
    "        self.run.logger = partial(self.mbar.write, table=True)\n",
    "        \n",
    "    def after_fit(self): self.mbar.on_iter_end()\n",
    "    def after_batch(self): self.pb.update(self.iter)\n",
    "    def begin_epoch   (self): self.set_pb()\n",
    "    def begin_validate(self): self.set_pb()\n",
    "        \n",
    "    def set_pb(self):\n",
    "        self.pb = progress_bar(self.dl, parent=self.mbar)#, auto_update=False)\n",
    "        self.mbar.update(self.epoch)\n",
    "\n",
    "#export\n",
    "################################## CUDA  #############################\n",
    "class CudaCallback(Callback):\n",
    "    def __init__(self,device): self.device=device \n",
    "    def begin_fit(self): self.model.to(self.device) \n",
    "    def begin_batch(self): self.run.xb,self.run.yb = self.xb.to(self.device),self.yb.to(self.device)            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StatefulOptimizer(Optimizer):\n",
    "    def __init__(self, params, steppers, stats=None, **defaults):\n",
    "        self.stats = listify(stats)\n",
    "        maybe_update(self.stats, defaults, get_defaults)\n",
    "        super().__init__(params, steppers, **defaults)\n",
    "        self.state = {}\n",
    "\n",
    "    def step(self):\n",
    "        for p,hyper in self.grad_params():\n",
    "            if p not in self.state:\n",
    "                self.state[p] = {}\n",
    "                maybe_update(self.stats, self.state[p], lambda o: o.init_state(p))\n",
    "            state = self.state[p]\n",
    "            for stat in self.stats: state = stat.update(p, state, **hyper)\n",
    "            compose(p, self.steppers, **state, **hyper)\n",
    "            self.state[p] = state\n",
    "            \n",
    "class Stat():\n",
    "    _defaults = {}\n",
    "    def init_state(self, p): raise NotImplementedError\n",
    "    def update(self, p, state, **kwargs): raise NotImplementedError\n",
    "\n",
    "def momentum_step(p, lr, grad_avg, **kwargs):\n",
    "    p.data.add_(-lr, grad_avg)\n",
    "    return p\n",
    "\n",
    "def lin_comb(v1, v2, beta): return beta*v1 + (1-beta)*v2\n",
    "\n",
    "\n",
    "class AverageGrad(Stat):\n",
    "    _defaults = dict(mom=0.9)\n",
    "    \n",
    "    def __init__(self, dampening:bool=False): self.dampening=dampening\n",
    "    def init_state(self, p): return {'grad_avg': torch.zeros_like(p.grad.data)}\n",
    "    def update(self, p, state, mom, **kwargs):\n",
    "        state['mom_damp'] = 1-mom if self.dampening else 1.\n",
    "        state['grad_avg'].mul_(mom).add_(state['mom_damp'], p.grad.data)\n",
    "        return state\n",
    "    \n",
    "class AverageSqrGrad(Stat):\n",
    "    _defaults = dict(sqr_mom=0.99)\n",
    "    \n",
    "    def __init__(self, dampening:bool=True): self.dampening=dampening\n",
    "    def init_state(self, p): return {'sqr_avg': torch.zeros_like(p.grad.data)}\n",
    "    def update(self, p, state, sqr_mom, **kwargs):\n",
    "        state['sqr_damp'] = 1-sqr_mom if self.dampening else 1.\n",
    "        state['sqr_avg'].mul_(sqr_mom).addcmul_(state['sqr_damp'], p.grad.data, p.grad.data)\n",
    "        return state\n",
    "\n",
    "class StepCount(Stat):\n",
    "    def init_state(self, p): return {'step': 0}\n",
    "    def update(self, p, state, **kwargs):\n",
    "        state['step'] += 1\n",
    "        return state\n",
    "    \n",
    "def debias(mom, damp, step): return damp * (1 - mom**step) / (1-mom)\n",
    "\n",
    "def adam_step(p, lr, mom, mom_damp, step, sqr_mom, sqr_damp, grad_avg, sqr_avg, eps, **kwargs):\n",
    "    debias1 = debias(mom,     mom_damp, step)\n",
    "    debias2 = debias(sqr_mom, sqr_damp, step)\n",
    "    p.data.addcdiv_(-lr / debias1, grad_avg, (sqr_avg/debias2 ).sqrt()  + eps )\n",
    "    return p\n",
    "adam_step._defaults = dict(eps=1e-5)\n",
    "\n",
    "def adam_opt(xtra_step=None, **kwargs):\n",
    "    return partial(StatefulOptimizer, steppers=[adam_step,weight_decay]+listify(xtra_step),\n",
    "                   stats=[AverageGrad(dampening=True), AverageSqrGrad(), StepCount()], **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def jax_get_learner(data, layers, loss_func, cb_funcs, opt_func,lr ):\n",
    "    model = nn.Sequential(*layers)\n",
    "    return Learner(model, data, loss_func, cb_funcs, opt_func,lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "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='2', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      \n",
       "    </div>\n",
       "    \n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>train_accuracy</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>valid_accuracy</th>\n",
       "      <th>time</th>\n",
       "      <th>lr</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.016902</td>\n",
       "      <td>0.572451</td>\n",
       "      <td>0.007822</td>\n",
       "      <td>0.699093</td>\n",
       "      <td>1:56:28</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.005200</td>\n",
       "      <td>0.773527</td>\n",
       "      <td>0.008704</td>\n",
       "      <td>0.719957</td>\n",
       "      <td>05:06</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 52min 47s, sys: 3min 23s, total: 56min 11s\n",
      "Wall time: 2h 1min 35s\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "bs = 64\n",
    "train_dataset = jax_Dataset(root_dir='./train')\n",
    "train_split_ds, valid_split_ds = torch.utils.data.random_split(train_dataset,\n",
    "                    [round(train_dataset.size*0.8), train_dataset.size-round(train_dataset.size*0.8)])\n",
    "train_dl,valid_dl = get_dls(train_split_ds, valid_split_ds, bs,num_workers=4,drop_last=True)\n",
    "data = DataBunch(*get_dls(train_split_ds, valid_split_ds, bs))\n",
    "cb_funcs = [ProgressCallback,partial(AvgStatsCallback,[accuracy]),partial(CudaCallback,device)]\n",
    "adam_opt_func = adam_opt()\n",
    "learn = jax_get_learner(data=data,layers=get_model_layers_3d(),loss_func=nn.MSELoss(),cb_funcs= cb_funcs, opt_func=adam_opt_func, lr=0.001)\n",
    "model_init(learn.model)\n",
    "\n",
    "%time learn.fit(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "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='2', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      \n",
       "    </div>\n",
       "    \n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>train_accuracy</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>valid_accuracy</th>\n",
       "      <th>time</th>\n",
       "      <th>lr</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.003195</td>\n",
       "      <td>0.844249</td>\n",
       "      <td>0.009021</td>\n",
       "      <td>0.773055</td>\n",
       "      <td>05:09</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.002268</td>\n",
       "      <td>0.885854</td>\n",
       "      <td>0.003086</td>\n",
       "      <td>0.861743</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 21min 7s, sys: 1min 11s, total: 22min 18s\n",
      "Wall time: 10min 20s\n"
     ]
    }
   ],
   "source": [
    "%time learn.fit(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "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='40', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      \n",
       "    </div>\n",
       "    \n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>train_accuracy</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>valid_accuracy</th>\n",
       "      <th>time</th>\n",
       "      <th>lr</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.017486</td>\n",
       "      <td>0.572059</td>\n",
       "      <td>0.011356</td>\n",
       "      <td>0.660030</td>\n",
       "      <td>05:35</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.005283</td>\n",
       "      <td>0.773031</td>\n",
       "      <td>0.003534</td>\n",
       "      <td>0.830089</td>\n",
       "      <td>05:04</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.003189</td>\n",
       "      <td>0.844881</td>\n",
       "      <td>0.002661</td>\n",
       "      <td>0.872505</td>\n",
       "      <td>04:56</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.002335</td>\n",
       "      <td>0.883372</td>\n",
       "      <td>0.001756</td>\n",
       "      <td>0.913533</td>\n",
       "      <td>05:06</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.001761</td>\n",
       "      <td>0.912328</td>\n",
       "      <td>0.001624</td>\n",
       "      <td>0.926338</td>\n",
       "      <td>05:11</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.001382</td>\n",
       "      <td>0.933256</td>\n",
       "      <td>0.001269</td>\n",
       "      <td>0.941084</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.001152</td>\n",
       "      <td>0.947459</td>\n",
       "      <td>0.001066</td>\n",
       "      <td>0.953630</td>\n",
       "      <td>05:13</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.000949</td>\n",
       "      <td>0.959971</td>\n",
       "      <td>0.001149</td>\n",
       "      <td>0.949057</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.000829</td>\n",
       "      <td>0.968012</td>\n",
       "      <td>0.000917</td>\n",
       "      <td>0.963374</td>\n",
       "      <td>05:09</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.000763</td>\n",
       "      <td>0.972704</td>\n",
       "      <td>0.000825</td>\n",
       "      <td>0.967539</td>\n",
       "      <td>05:04</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.000699</td>\n",
       "      <td>0.977020</td>\n",
       "      <td>0.000690</td>\n",
       "      <td>0.976515</td>\n",
       "      <td>05:11</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.000641</td>\n",
       "      <td>0.980687</td>\n",
       "      <td>0.000680</td>\n",
       "      <td>0.976520</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.000595</td>\n",
       "      <td>0.983561</td>\n",
       "      <td>0.000606</td>\n",
       "      <td>0.981279</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>0.000559</td>\n",
       "      <td>0.985776</td>\n",
       "      <td>0.000575</td>\n",
       "      <td>0.983299</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>0.000532</td>\n",
       "      <td>0.987541</td>\n",
       "      <td>0.000559</td>\n",
       "      <td>0.984062</td>\n",
       "      <td>05:08</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>0.000505</td>\n",
       "      <td>0.989005</td>\n",
       "      <td>0.000666</td>\n",
       "      <td>0.978558</td>\n",
       "      <td>05:08</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>16</td>\n",
       "      <td>0.000478</td>\n",
       "      <td>0.990388</td>\n",
       "      <td>0.000528</td>\n",
       "      <td>0.985249</td>\n",
       "      <td>05:08</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17</td>\n",
       "      <td>0.000457</td>\n",
       "      <td>0.991532</td>\n",
       "      <td>0.000546</td>\n",
       "      <td>0.984681</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>18</td>\n",
       "      <td>0.000446</td>\n",
       "      <td>0.992165</td>\n",
       "      <td>0.000528</td>\n",
       "      <td>0.986035</td>\n",
       "      <td>05:12</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>19</td>\n",
       "      <td>0.000424</td>\n",
       "      <td>0.993119</td>\n",
       "      <td>0.000509</td>\n",
       "      <td>0.987013</td>\n",
       "      <td>05:06</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>0.000410</td>\n",
       "      <td>0.993927</td>\n",
       "      <td>0.000512</td>\n",
       "      <td>0.986940</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>21</td>\n",
       "      <td>0.000395</td>\n",
       "      <td>0.994590</td>\n",
       "      <td>0.000485</td>\n",
       "      <td>0.988167</td>\n",
       "      <td>05:09</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>22</td>\n",
       "      <td>0.000385</td>\n",
       "      <td>0.994981</td>\n",
       "      <td>0.000485</td>\n",
       "      <td>0.987642</td>\n",
       "      <td>05:11</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>23</td>\n",
       "      <td>0.000374</td>\n",
       "      <td>0.995402</td>\n",
       "      <td>0.000456</td>\n",
       "      <td>0.989561</td>\n",
       "      <td>05:11</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24</td>\n",
       "      <td>0.000361</td>\n",
       "      <td>0.995965</td>\n",
       "      <td>0.000486</td>\n",
       "      <td>0.987914</td>\n",
       "      <td>05:13</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>25</td>\n",
       "      <td>0.000353</td>\n",
       "      <td>0.996312</td>\n",
       "      <td>0.000433</td>\n",
       "      <td>0.990102</td>\n",
       "      <td>05:12</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>26</td>\n",
       "      <td>0.000341</td>\n",
       "      <td>0.996642</td>\n",
       "      <td>0.000478</td>\n",
       "      <td>0.989123</td>\n",
       "      <td>05:14</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>27</td>\n",
       "      <td>0.000333</td>\n",
       "      <td>0.996962</td>\n",
       "      <td>0.000454</td>\n",
       "      <td>0.989767</td>\n",
       "      <td>05:14</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>28</td>\n",
       "      <td>0.000324</td>\n",
       "      <td>0.997271</td>\n",
       "      <td>0.000421</td>\n",
       "      <td>0.990907</td>\n",
       "      <td>05:10</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>29</td>\n",
       "      <td>0.000320</td>\n",
       "      <td>0.997358</td>\n",
       "      <td>0.000433</td>\n",
       "      <td>0.990177</td>\n",
       "      <td>05:13</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>0.000316</td>\n",
       "      <td>0.997437</td>\n",
       "      <td>0.000450</td>\n",
       "      <td>0.989482</td>\n",
       "      <td>05:09</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>31</td>\n",
       "      <td>0.000305</td>\n",
       "      <td>0.997853</td>\n",
       "      <td>0.000432</td>\n",
       "      <td>0.990018</td>\n",
       "      <td>05:13</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>32</td>\n",
       "      <td>0.000298</td>\n",
       "      <td>0.997981</td>\n",
       "      <td>0.000395</td>\n",
       "      <td>0.991732</td>\n",
       "      <td>05:12</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>33</td>\n",
       "      <td>0.000293</td>\n",
       "      <td>0.998105</td>\n",
       "      <td>0.000386</td>\n",
       "      <td>0.991995</td>\n",
       "      <td>05:12</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>34</td>\n",
       "      <td>0.000286</td>\n",
       "      <td>0.998297</td>\n",
       "      <td>0.000417</td>\n",
       "      <td>0.990827</td>\n",
       "      <td>05:12</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>35</td>\n",
       "      <td>0.000280</td>\n",
       "      <td>0.998394</td>\n",
       "      <td>0.000400</td>\n",
       "      <td>0.991538</td>\n",
       "      <td>05:15</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>36</td>\n",
       "      <td>0.000276</td>\n",
       "      <td>0.998495</td>\n",
       "      <td>0.000378</td>\n",
       "      <td>0.992643</td>\n",
       "      <td>05:06</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>37</td>\n",
       "      <td>0.000271</td>\n",
       "      <td>0.998614</td>\n",
       "      <td>0.000384</td>\n",
       "      <td>0.992327</td>\n",
       "      <td>05:11</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>38</td>\n",
       "      <td>0.000270</td>\n",
       "      <td>0.998606</td>\n",
       "      <td>0.000378</td>\n",
       "      <td>0.992357</td>\n",
       "      <td>05:11</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>39</td>\n",
       "      <td>0.000263</td>\n",
       "      <td>0.998772</td>\n",
       "      <td>0.000395</td>\n",
       "      <td>0.992200</td>\n",
       "      <td>05:12</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7h 6min 40s, sys: 21min 31s, total: 7h 28min 12s\n",
      "Wall time: 3h 27min 20s\n"
     ]
    }
   ],
   "source": [
    "%time learn.fit(40)\n",
    "# 保存\n",
    "torch.save(learn.model, './0922.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9311, device='cuda:0')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "######################################### testing ####################################################\n",
    "device = 'cuda'\n",
    "bs = 64\n",
    "acc_threshold = 1/22\n",
    "test_dataset = jax_Dataset(root_dir='./test')\n",
    "test_dl = DataLoader( test_dataset, batch_size=bs, shuffle=False)\n",
    "learn.in_train = False\n",
    "accTmp = 0\n",
    "accTot = 0\n",
    "accCnt = 0\n",
    "accMean = 0\n",
    "for i,(xb,yb) in enumerate(test_dl): \n",
    "    xb = xb.unsqueeze(1)           \n",
    "    xb = (xb.to(device)).float()   \n",
    "    yb = (yb.to(device)).float()   \n",
    "    yh = learn.model(xb)\n",
    "    accTmp = accuracy(yh,yb)\n",
    "    accTot += accTmp\n",
    "    accCnt += 1\n",
    "    \n",
    "accMean = accTot/accCnt\n",
    "accMean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9769, device='cuda:0')"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "######################################### testing ####################################################\n",
    "device = 'cuda'\n",
    "bs = 64\n",
    "acc_threshold = 3/44\n",
    "test_dataset = jax_Dataset(root_dir='./test')\n",
    "test_dl = DataLoader( test_dataset, batch_size=bs, shuffle=False)\n",
    "learn.in_train = False\n",
    "accTmp = 0\n",
    "accTot = 0\n",
    "accCnt = 0\n",
    "accMean = 0\n",
    "\n",
    "for i,(xb,yb) in enumerate(test_dl): \n",
    "    xb = xb.unsqueeze(1)           \n",
    "    xb = (xb.to(device)).float()   \n",
    "    yb = (yb.to(device)).float()   \n",
    "    yh = learn.model(xb)\n",
    "    accTmp = accuracy(yh,yb)\n",
    "    accTot += accTmp\n",
    "    accCnt += 1\n",
    "\n",
    "accMean = accTot/accCnt\n",
    "accMean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "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='20', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      \n",
       "    </div>\n",
       "    \n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>train_accuracy</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>valid_accuracy</th>\n",
       "      <th>time</th>\n",
       "      <th>lr</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.000258</td>\n",
       "      <td>0.998887</td>\n",
       "      <td>0.000388</td>\n",
       "      <td>0.991784</td>\n",
       "      <td>05:13</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.000254</td>\n",
       "      <td>0.998901</td>\n",
       "      <td>0.000375</td>\n",
       "      <td>0.992650</td>\n",
       "      <td>05:15</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.000250</td>\n",
       "      <td>0.999012</td>\n",
       "      <td>0.000388</td>\n",
       "      <td>0.992194</td>\n",
       "      <td>05:20</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.000248</td>\n",
       "      <td>0.999010</td>\n",
       "      <td>0.000372</td>\n",
       "      <td>0.992495</td>\n",
       "      <td>05:22</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.000246</td>\n",
       "      <td>0.999072</td>\n",
       "      <td>0.000346</td>\n",
       "      <td>0.993563</td>\n",
       "      <td>06:04</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.000240</td>\n",
       "      <td>0.999171</td>\n",
       "      <td>0.000376</td>\n",
       "      <td>0.993159</td>\n",
       "      <td>06:04</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.000237</td>\n",
       "      <td>0.999162</td>\n",
       "      <td>0.000352</td>\n",
       "      <td>0.993658</td>\n",
       "      <td>05:27</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.000233</td>\n",
       "      <td>0.999271</td>\n",
       "      <td>0.000357</td>\n",
       "      <td>0.993282</td>\n",
       "      <td>05:24</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.000231</td>\n",
       "      <td>0.999279</td>\n",
       "      <td>0.000345</td>\n",
       "      <td>0.993684</td>\n",
       "      <td>05:17</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.000228</td>\n",
       "      <td>0.999301</td>\n",
       "      <td>0.000338</td>\n",
       "      <td>0.993836</td>\n",
       "      <td>05:13</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.000226</td>\n",
       "      <td>0.999351</td>\n",
       "      <td>0.000346</td>\n",
       "      <td>0.993695</td>\n",
       "      <td>05:26</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>0.000222</td>\n",
       "      <td>0.999399</td>\n",
       "      <td>0.000333</td>\n",
       "      <td>0.993984</td>\n",
       "      <td>05:29</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>0.000220</td>\n",
       "      <td>0.999408</td>\n",
       "      <td>0.000350</td>\n",
       "      <td>0.993462</td>\n",
       "      <td>05:21</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>0.000217</td>\n",
       "      <td>0.999433</td>\n",
       "      <td>0.000337</td>\n",
       "      <td>0.993925</td>\n",
       "      <td>05:24</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>0.000213</td>\n",
       "      <td>0.999497</td>\n",
       "      <td>0.000335</td>\n",
       "      <td>0.994024</td>\n",
       "      <td>05:28</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>0.000215</td>\n",
       "      <td>0.999457</td>\n",
       "      <td>0.000350</td>\n",
       "      <td>0.993412</td>\n",
       "      <td>05:20</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>16</td>\n",
       "      <td>0.000209</td>\n",
       "      <td>0.999553</td>\n",
       "      <td>0.000332</td>\n",
       "      <td>0.994105</td>\n",
       "      <td>05:19</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17</td>\n",
       "      <td>0.000207</td>\n",
       "      <td>0.999556</td>\n",
       "      <td>0.000337</td>\n",
       "      <td>0.994046</td>\n",
       "      <td>05:26</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>18</td>\n",
       "      <td>0.000205</td>\n",
       "      <td>0.999524</td>\n",
       "      <td>0.000329</td>\n",
       "      <td>0.994222</td>\n",
       "      <td>05:22</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>19</td>\n",
       "      <td>0.000204</td>\n",
       "      <td>0.999575</td>\n",
       "      <td>0.000322</td>\n",
       "      <td>0.994365</td>\n",
       "      <td>05:46</td>\n",
       "      <td>0.001</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3h 44min 28s, sys: 11min 35s, total: 3h 56min 4s\n",
      "Wall time: 1h 49min 10s\n"
     ]
    }
   ],
   "source": [
    "%time learn.fit(20)\n",
    "# 保存\n",
    "torch.save(learn.model, './0922-2.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9813, device='cuda:0')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "######################################### testing ####################################################\n",
    "device = 'cuda'\n",
    "bs = 64\n",
    "acc_threshold = 3/44\n",
    "test_dataset = jax_Dataset(root_dir='./test')\n",
    "test_dl = DataLoader( test_dataset, batch_size=bs, shuffle=False)\n",
    "learn.in_train = False\n",
    "accTmp = 0\n",
    "accTot = 0\n",
    "accCnt = 0\n",
    "accMean = 0\n",
    "\n",
    "for i,(xb,yb) in enumerate(test_dl): \n",
    "    xb = xb.unsqueeze(1)           \n",
    "    xb = (xb.to(device)).float()   \n",
    "    yb = (yb.to(device)).float()   \n",
    "    yh = learn.model(xb)\n",
    "    accTmp = accuracy(yh,yb)\n",
    "    accTot += accTmp\n",
    "    accCnt += 1\n",
    "\n",
    "accMean = accTot/accCnt\n",
    "accMean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "## for test$$\n",
    "def paramRMSE(yh,yb):\n",
    "    tmp = torch.abs(yh-yb)\n",
    "    largpart = tmp.mul(tmp)\n",
    "    return largpart.sum().sqrt()\n",
    "def paramAccu(yh,yb,acc_threshold): \n",
    "    tmp = (torch.abs(yh-yb) <= acc_threshold).float()\n",
    "    return tmp.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "############## testing ###########\n",
    "## save ***.csv and ****.npy and comExp.csv\n",
    "fpSave_rslt_csv = './test/rslt.csv';\n",
    "\n",
    "f = open(fpSave_rslt_csv,'w',newline='',encoding='UTF-8')\n",
    "\n",
    "csv_writer = csv.writer(f)\n",
    "# csv_writer.writerow([\"xO\",\"yO\",\"zO\",\"pA1\",\"pR1\",\"pA2\",\"pR2\",\"TM\"])\n",
    "#csv_cmp_writer do not write head\n",
    "######################################### testing ####################################################\n",
    "device = 'cuda'\n",
    "acc_threshold = 3/44\n",
    "cb_funcs = [ProgressCallback,partial(AvgStatsCallback,[accuracy]),partial(CudaCallback,device)]\n",
    "adam_opt_func = adam_opt()\n",
    "\n",
    "test_bs = 1\n",
    "test_dataset = jax_Dataset(root_dir='./test')\n",
    "test_dl = DataLoader( test_dataset, batch_size=test_bs, shuffle=False)\n",
    "learn = jax_get_learner(data=None,layers=get_model_layers_3d(),loss_func=nn.MSELoss(),cb_funcs= cb_funcs, opt_func=adam_opt_func, lr=0.001)\n",
    "learn.in_train = False\n",
    "learn.model = torch.load('./0922-2.pkl')\n",
    "accTmp = 0\n",
    "accTot = 0\n",
    "accCnt = 0\n",
    "accMean = 0\n",
    "testTime = 0\n",
    "paramacc1 = 0\n",
    "paramrmse1 = 0\n",
    "paramacc2 = 0\n",
    "paramrmse2 = 0\n",
    "\n",
    "for i,(xb,yb) in enumerate(test_dl): \n",
    "    xb = xb.unsqueeze(1)           \n",
    "    xb = (xb.to(device)).float()   \n",
    "    yb = (yb.to(device)).float()   \n",
    "    \n",
    "    Tstart = time.time()\n",
    "    yh = learn.model(xb)\n",
    "    Tend = time.time()\n",
    "    testTime = Tend-Tstart\n",
    "    \n",
    "    acc_threshold = 3/44 \n",
    "    paramacc1  = paramAccu(yh,yb,acc_threshold)\n",
    "    paramrmse1 = paramRMSE(yh,yb)\n",
    "    acc_threshold = 1/22\n",
    "    paramacc2  = paramAccu(yh,yb,acc_threshold)\n",
    "    paramrmse2 = paramRMSE(yh,yb)\n",
    "    csv_writer.writerow([(yh[0,0].to('cpu')).detach().numpy(), \n",
    "                         (yh[0,1].to('cpu')).detach().numpy(),\n",
    "                         (yh[0,2].to('cpu')).detach().numpy(), \n",
    "                         \n",
    "                         (yb[0,0].to('cpu')).detach().numpy(), \n",
    "                         (yb[0,1].to('cpu')).detach().numpy(),\n",
    "                         (yb[0,2].to('cpu')).detach().numpy(),\n",
    "                         \n",
    "                         paramacc1.to('cpu').detach().numpy(),\n",
    "                         paramrmse1.to('cpu').detach().numpy(),\n",
    "                         paramacc2.to('cpu').detach().numpy(),\n",
    "                         paramrmse2.to('cpu').detach().numpy(),\n",
    "                         \n",
    "                         testTime]);\n",
    "    accTot += paramacc1\n",
    "    accCnt += 1\n",
    "\n",
    "\n",
    "accMean = accTot/accCnt\n",
    "accMean\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}