tabular_experiment_model.py

from fastai import torch_core, layers
from fastai.basic_data import *
from fastai.core import *
from fastai.layers import embedding
from fastai.basic_train import Learner
from torch import nn, optim, as_tensor, Tensor
import torch
import logging

def tabularexperiment_learner(data:DataBunch, layers:Collection[int], emb_szs:Dict[str,int]=None, metrics=None,
        ps:Collection[float]=None, emb_drop:float=0., y_range:OptRange=None, use_bn:bool=True, **learn_kwargs):
    "Get a `Learner` using `data`, with `metrics`, including a `TabularModel` created using the remaining params."
    emb_szs = data.get_emb_szs(ifnone(emb_szs, {}))
    model = TabularExperimentModel(emb_szs, len(data.cont_names), out_sz=data.c, layers=layers, ps=ps, emb_drop=emb_drop,
                         y_range=y_range, use_bn=use_bn)
    learner = Learner(data, model, metrics=metrics, **learn_kwargs)
    learner.to_onnx = partial(to_onnx, data=data, model=model)
    return learner

def to_onnx(file_path:str, data, model:str):    
    logging.info('getting data...')
    dummy_input = next(iter(data.train_dl))[0]
    logging.info(f'saviing {file_path} file ...')
    torch.onnx.export(model, (dummy_input[0],dummy_input[1]), file_path)

class TabularExperimentModel (nn.Module):
    def __init__(self, emb_szs:ListSizes, n_cont:int, out_sz:int, layers:Collection[int], ps:Collection[float]=None,
                 emb_drop:float=0., y_range:OptRange=None, use_bn:bool=True, bn_final:bool=False):
        super().__init__()
        ps = ifnone(ps, [0]*len(layers))
        ps = listify(ps, layers)
        self.embeds = nn.ModuleList([embedding(ni, nf) for ni,nf in emb_szs])
        self.emb_drop = nn.Dropout(emb_drop)
        self.bn_cont = nn.BatchNorm1d(n_cont)
        #self.bn_sigm = nn.BatchNorm1d(out_sz)
        n_emb = sum(e.embedding_dim for e in self.embeds)
        self.n_emb,self.n_cont,self.y_range = n_emb,n_cont,y_range
        sizes = self._get_sizes(layers, out_sz)
        actns = [nn.ReLU(inplace=True)]*(len(sizes)-2)
        layers = []
        for i,(n_in,n_out,dp,act) in enumerate(zip(sizes[:-2],sizes[1:-1],[0.]+ps,actns)):
            layers += self.lin_bn_drop(n_in, n_out, bn=use_bn, p=dp, actn=act)
        layers.append(nn.Linear(sizes[-2], sizes[-1]))
        if bn_final: layers.append(nn.BatchNorm1d(sizes[-1]))
        self.layers = nn.Sequential(*layers)

    def lin_bn_drop(self, n_in:int, n_out:int, bn:bool=True, p:float=0., actn:Optional[nn.Module]=None):
        "Sequence of batchnorm (if `bn`), dropout (with `p`) and linear (`n_in`,`n_out`) layers followed by `actn`."
        layers = [nn.Linear(n_in, n_out)]
        if bn:
            layers.append(nn.BatchNorm1d(n_out))

        if p > 0: 
            layers.append(nn.Dropout(p))

        if actn is not None: 
            layers.append(actn)
        return layers

    def _get_sizes(self, layers, out_sz):
        return [self.n_emb + self.n_cont] + layers + [out_sz]

    def forward(self, x_cat:Tensor, x_cont:Tensor) -> Tensor:
        if self.n_emb > 0:
            x_emb = [e(x_cat[:,i]) for i,e in enumerate(self.embeds)]
            x_emb = torch.cat(x_emb, 1)
            x_feat = self.emb_drop(x_emb)
        if self.n_cont > 0:
            x_cont = self.bn_cont(x_cont)
            x_feat = torch.cat([x_feat, x_cont], 1) if self.n_emb != 0 else x_cont
        x_feat = self.layers(x_feat)
        if self.y_range is not None:
            #x_emb = self.bn_sigm(x_emb)
            x_feat = torch.sigmoid(x_feat)
            x_feat = (self.y_range[1]-self.y_range[0]) * x_feat + self.y_range[0]
        return x_feat