<jupyter_start><jupyter_text>N-Beats model> A basic architecture for time series forecasting.<jupyter_code># hide
from nbdev.showdoc import *
from fastcore.test import *
#export
from fastcore.utils import *
from fastcore.imports import *
from fastai.basics import *
from fastai.callback.hook import num_features_model
from fastai.callback.all import *
from fastai.torch_core import *
from torch.autograd import Variable
from fastseq.all import *<jupyter_output><empty_output><jupyter_text>Blocks<jupyter_code># export

def linspace(lookback, horizon):
    lin_space = torch.linspace(
        -lookback, horizon, lookback + horizon
    )
    b_ls = lin_space[:lookback]
    f_ls = lin_space[lookback:]
    return b_ls, f_ls

# export
class Block(Module):
    def __init__(self, fnc_f, fnc_b=None):
        sizes = [self.lookback] + self.layers
        ps = ifnone(self.ps, L([0])*len(self.layers))
        actns = [nn.ReLU(inplace=True) for _ in range(len(sizes)-2)] + [None]
        _layers = [LinBnDrop(sizes[i], sizes[i+1], bn=self.use_bn, p=p, act=a)
                       for i,(p,a) in enumerate(zip(ps, actns))]
        self.layers = nn.Sequential(*_layers)
        self.att = LinBnDrop(sizes[-1], 1)
        if self.share_thetas:
            self.theta_f_fc = self.theta_b_fc = LinBnDrop(sizes[-1], self.thetas_dim)
        else:
            self.theta_b_fc = LinBnDrop(sizes[-1], self.thetas_dim)
            self.theta_f_fc = LinBnDrop(sizes[-1], self.thetas_dim)
            
        b, f = linspace(self.lookback, self.horizon)
        self.backcast_linspace = Variable(b, requires_grad=False).to(self.device)
        self.forecast_linspace = Variable(f, requires_grad=False).to(self.device)
        self.fnc_f = fnc_f
        self.fnc_b = ifnone(fnc_b, fnc_f)
        self.to(self.device)
        self.y_range = getattr(self,'y_range', None)
        
    def forward(self, x):
        # trend
        x = self.layers(x)
        att = torch.sigmoid(self.att(x))
        theta_b = self.apply_range(self.theta_b_fc(x)) * att
        theta_f = self.apply_range(self.theta_f_fc(x)) * att
        backcast = self.fnc_b(theta_b, self.backcast_linspace)
        forecast = self.fnc_f(theta_f, self.forecast_linspace)
        return {'b':backcast,'f': forecast, 'theta': theta_b + theta_f}
    
    def apply_range(self, x): 
        if self.y_range is None:
            return x
        return (self.y_range[1]-self.y_range[0]) * torch.sigmoid(x) + self.y_range[0]
# export
def bias_model(thetas, t):
    r= torch.mm(t[None,:].float().T,thetas[:,0][None,:]).T 
    if thetas.shape[-1]==2:
        return r+thetas[:,1][:,None]
    return r

class BiasBlock(Block):
    def __init__(
        self, layers:L, device, thetas_dim=2, lookback=10, horizon=5, use_bn=True, bn_final=False, ps:L=None
    ):
        share_thetas=True
        assert thetas_dim <= 2, f"thetas_dim for BaisBlock must be < than 2, is now {thetas_dim}"
        store_attr("device,layers,thetas_dim,use_bn,ps,lookback,horizon,bn_final,share_thetas" ) 
        self.layers=L(self.layers[-1])
        super().__init__(bias_model)        
        self.to(device)
mdl = BiasBlock([256], device=torch.device('cuda'), lookback=100, horizon=10)
dct = mdl(torch.randn(64,100).to(torch.device('cuda')))
test_eq(dct['b'].shape, (64,100))
test_eq(dct['f'].shape, (64,10))
test_eq(dct['theta'].shape,(64,2))
# export
class LinearD(nn.Linear):        
    def forward(self, x, *args, **kwargs):
        return super().forward(x)
        
class GenericBlock(Block):
    def __init__(
        self, layers:L, thetas_dim:int, device, lookback=10, horizon=5, use_bn=True, bn_final=False, ps:L=None, share_thetas=True, y_range=[-.5,.5]
    ):
        store_attr("y_range,device,layers,thetas_dim,use_bn,ps,lookback,horizon,bn_final,share_thetas" ) 
        super().__init__(LinearD(self.thetas_dim, self.horizon),LinearD(self.thetas_dim, self.lookback))        
        self.to(device)
        
mdl = GenericBlock([256], thetas_dim=4, device=torch.device('cuda'), lookback=100, horizon=10)
dct = mdl(torch.randn(64,100).to(torch.device('cuda')))
test_eq(dct['b'].shape, (64,100))
test_eq(dct['f'].shape, (64,10))
# export

def seasonality_model(thetas, t):
    p = thetas.size()[-1]
    assert p < 10, "thetas_dim is too big."
    p1, p2 = (p // 2, p // 2) if p % 2 == 0 else (p // 2, p // 2 + 1)
    s1 = [torch.cos(2 * np.pi * i * t)[None,:] for i in range(p1)] # H/2-1
    s2 = [torch.sin(2 * np.pi * i * t)[None,:] for i in range(p2)]
    S = torch.cat([*s1, *s2])
    return thetas.mm(S)

class SeasonalityBlock(Block):
    def __init__(
        self, layers:L, thetas_dim:int, device, lookback=10, horizon=5, use_bn=True, bn_final=False, ps:L=None, share_thetas=False,y_range=[-.5,.5]
    ):       
        store_attr("y_range,device,layers,thetas_dim,use_bn,ps,lookback,horizon,bn_final,share_thetas" ) 
        super().__init__(seasonality_model )        
        self.to(device)
mdl = SeasonalityBlock([256], thetas_dim=4, device=torch.device('cuda'), lookback=100, horizon=10)
dct = mdl(torch.randn(64,100).to(torch.device('cuda')))
test_eq(dct['b'].shape, (64,100))
test_eq(dct['f'].shape, (64,10))
# export
def trend_model(thetas, t):
    p = thetas.size()[-1]
    assert p <= 4, "thetas_dim is too big."
    a = [torch.pow(t, i)[None,:] for i in range(p)]
    T = torch.cat(a).float()
    return thetas.mm(T)

class TrendBlock(Block):
    def __init__(
        self, layers:L, device, thetas_dim, lookback=10, horizon=5, use_bn=True, bn_final=False, ps:L=None, share_thetas=False, y_range=[-.05,.05]
    ):
        store_attr("y_range,device,layers,thetas_dim,use_bn,ps,lookback,horizon,bn_final,share_thetas" ) 
        super().__init__(trend_model)        
        self.to(device)
mdl = TrendBlock([256], thetas_dim=4, device=torch.device('cuda'), lookback=100, horizon=10)
dct = mdl(torch.randn(64,100).to(torch.device('cuda')))
test_eq(dct['b'].shape, (64,100))
test_eq(dct['f'].shape, (64,10))
test_eq(dct['theta'].shape,(64,4))
# export

# blocks= defaultdict(GenericBlock)
# blocks.update({"seasonality": SeasonalityBlock,
#          "trend": TrendBlock,
#          "generic": GenericBlock})
# blocks_int= {i:v for i,v in enumerate(blocks.items())}
# not pritty but still works better
def select_block(o):
    if isinstance(o,int):
        if o == 0:
            return SeasonalityBlock
        elif o == 1:
            return TrendBlock
        elif o == 2:
            return BaisBlock
        else:
            return GenericBlock
    else:
        if o == 'seasonality':
            return SeasonalityBlock
        elif o == 'trend':
            return TrendBlock
        elif o =='bias':
            return BiasBlock
        else:
            return GenericBlock
# export
class NBeatsNet(Module):
    def __init__(
        self,
        device,
        stack_types=('bias', 'trend', 'seaonality'),
        nb_blocks_per_stack=3,
        horizon=5,
        lookback=10,
        thetas_dim=None,
        share_weights_in_stack=False,
        layers= [200,100],
    ):
        super(NBeatsNet, self).__init__()
        thetas_dim = ifnone(thetas_dim,[3 if 'bias' not in o else 2 for o in stack_types  ])
        stack_types= L(stack_types)
        store_attr('device,horizon,lookback,layers,nb_blocks_per_stack,share_weights_in_stack,stack_types,thetas_dim,device')
        self.stacks = []
        self._str = "| N-Beats\n"
        
        self.bn = BatchNorm(lookback, ndim=2)        
        stacks = OrderedDict()        
        for stack_id in range(len(self.stack_types)):
            
            stacks[str(self.stack_types[stack_id]) + str(stack_id)] = self.create_stack(stack_id)
        self.stacks = nn.Sequential(stacks)
        
    def create_stack(self, stack_id):
        stack_type = self.stack_types[stack_id]
        self._str += f"| --  Stack {stack_type.title()} (#{stack_id}) (share_weights_in_stack={self.share_weights_in_stack})\n"
        
        blocks = []
        for block_id in range(self.nb_blocks_per_stack):
            block_init = select_block(stack_type)
            if self.share_weights_in_stack and block_id != 0:
                block = blocks[-1]  # pick up the last one when we share weights.
            else:
                block = block_init(
                    layers = self.layers,
                    thetas_dim = self.thetas_dim[stack_id],
                    device = self.device,
                    lookback = self.lookback,
                    horizon = self.horizon,
                )
            self._str += f"     | -- {block}\n"
            blocks.append(block)
            
        return nn.Sequential(*blocks)

    def forward(self, backcast):
        backcast = backcast.view([-1,backcast.shape[-1]])
        forecast = torch.zeros(
            size=(backcast.size()[0], self.horizon,)
        )  # maybe batch size here.
        
        dct = defaultdict(dict)
        for stack_id, names in enumerate(self.stacks.named_children()):
            name = names[0]
            for block_id in range(len(self.stacks[stack_id])):
                dct[name+'_'+str(block_id)] = self.stacks[stack_id][block_id](backcast)
                backcast = backcast.to(self.device) - dct[name+'_'+str(block_id)]['b']
                forecast = forecast.to(self.device) + dct[name+'_'+str(block_id)]['f']
        return forecast[:,None,:], backcast[:,None,:], dct


mdl = NBeatsNet(torch.device('cpu'), stack_types=('seaonality'), )
f, b, dct = mdl(torch.randn(64,1,10))
print(f.shape,b.shape)
print([(stack, o['theta'].mean(0)) for stack,o in dct.items()])<jupyter_output>torch.Size([64, 1, 5]) torch.Size([64, 1, 10])
[('seaonality0_0', tensor([-0.0027,  0.0063, -0.0135], grad_fn=<MeanBackward1>)), ('seaonality0_1', tensor([ 0.0082, -0.0095,  0.0045], grad_fn=<MeanBackward1>)), ('seaonality0_2', tensor([0.0034, 0.0038, 0.0144], grad_fn=<MeanBackward1>))]<jupyter_text>Metrics<jupyter_code># export
from fastai.data.all import *
from fastai.optimizer import *
from fastai.learner import *
#export
class NBeatsTheta(Metric):
    "The sqaure of the `theta` for every block. "
    def reset(self):           self.total,self.count = 0.,0
    def accumulate(self, learn):
        bs = find_bs(learn.yb)        
        self.total += to_detach(learn.n_beats_trainer.metrics['theta'].mean())*bs
        self.count += bs
    @property
    def value(self): return self.total/self.count if self.count != 0 else None
    @property
    def name(self):  return "theta"
#export
class NBeatsBackwards(Metric):
    "The loss according to the `loss_func` on the backwards part of the time-serie."
    def reset(self):           self.total,self.count = 0.,0
    def accumulate(self, learn):
        bs = find_bs(learn.yb)        
        self.total += to_detach(learn.n_beats_trainer.metrics['b_loss'].mean())*bs
        self.count += bs
    @property
    def value(self): return self.total/self.count if self.count != 0 else None
    @property
    def name(self):  return "b_loss"<jupyter_output><empty_output><jupyter_text>Callback<jupyter_code># export
class NBeatsTrainer(Callback):
    "`Callback` that adds weights regularization the thetas in N-Beats training."
    def __init__(self, theta=0., b_loss=0.): 
        self.theta, self.b_loss = theta, b_loss
        self.metrics = {'theta':tensor([0.]), 'b_loss':tensor([0.])}
        self.b = None

    def begin_train(self): 
        self.out = defaultdict(dict)
        self.metrics = {'theta':tensor([0.]), 'b_loss':tensor([0.])}
        
    def begin_validate(self): 
        self.out = defaultdict(dict)
        self.metrics = {'theta':tensor([0.]), 'b_loss':tensor([0.])}
        
    def after_pred(self):
        self.b = self.pred[1] 
        self.pred[2]['total_b'] = self.pred[1] 
        self.out = concat_dct(self.pred[2], self.out)   
        self.learn.pred = self.pred[0]

    def after_loss(self):        
        # theta
        value=tensor([0.])
        for key in self.out.keys():
            if 'bias' not in key and 'total' not in key:
                v = self.out[key]['theta'].float().pow(2).mean()
                if self.theta != 0.:     
                    self.learn.loss += self.theta * v.item()
                value = value + v
        self.metrics['theta'] += value.clone().cpu().detach()
        
        # backwards 
        value = self.learn.loss_func(self.b.float(), *self.xb, reduction='mean') 
        if self.b_loss != 0.:
            self.learn.loss += self.b_loss * value.mean() 
        self.metrics['b_loss'] += value.sum().clone().detach()
            
    def theta_means(self, df=True):
        ret ={key:self.out[key]['theta'].float().mean(0).cpu().data for key in self.out.keys() if 'total' not in key}
        if df:
            return pd.DataFrame(ret)
        return ret
# export
def CombinedLoss(*losses, ratio:dict=None):
    _ratio = defaultdict(lambda:1.)
    if ratio is not None:
        _ratio.update(ratio)    
    ratio = _ratio
    
    def _inner(pred, truth, *args,**kwargs):
        loss = None
        for _loss in losses:
            if loss is None:
                loss = ratio[_loss.__name__] * _loss(pred, truth, *args,**kwargs)
            else:
                loss += ratio[_loss.__name__] * _loss(pred, truth, *args,**kwargs)
        return loss
    
    return _inner
#     metrics = L()
#     for _loss in losses:
#         metrics.append(lambda pred, truth, *args,**kwargs: ratio[_loss.__name__] * _loss(pred, truth, *args,**kwargs))
#     print(metrics)
  
y, y_hat = torch.arange(10).float(), torch.arange(10).float()+torch.randn(10)
loss_fnc = CombinedLoss(F.mse_loss,smape)
test_eq(F.mse_loss(y,y_hat)+smape(y,y_hat),loss_fnc(y, y_hat))

r = 10
loss_fnc = CombinedLoss(F.mse_loss, smape, ratio = {'mse_loss':r})
test_eq(r*F.mse_loss(y,y_hat)+smape(y,y_hat),loss_fnc(y, y_hat))<jupyter_output><empty_output><jupyter_text>Learner<jupyter_code># export
from fastai.basics import *
from fastseq.all import *

@delegates(NBeatsNet.__init__)
def nbeats_learner(dbunch:TSDataLoaders, output_channels=None, metrics=None,cbs=None, theta=0., b_loss=0., loss_func=None, **kwargs):
    "Build a N-Beats style learner"    
    model = NBeatsNet(
        device = dbunch.train_dl.device,
        horizon = dbunch.train_dl.horizon,
        lookback = dbunch.train_dl.lookback, 
        **kwargs
       )
    
    loss_func = ifnone(loss_func, CombinedLoss(F.mse_loss, smape, ratio = {'smape':.0005}))
    learn = Learner(dbunch, model, loss_func=loss_func, opt_func= Adam, 
                    metrics=L(metrics)+L(mae, smape, F.mse_loss, NBeatsTheta(), NBeatsBackwards()),
                    cbs=L(NBeatsTrainer(theta, b_loss))+L(cbs)
                   )
    learn.lh = (learn.dbunch.train_dl.lookback/learn.dbunch.train_dl.horizon)
    return learn
# hide
from fastseq.data.external import *
items = dummy_data_generator(50, 10, nrows=100)
data = TSDataLoaders.from_items(items, horizon = 7)
learn = nbeats_learner(data,layers=[10], theta = .1, b_loss=.3, stack_types=("seasonality","trend"),
                       loss_func= CombinedLoss(F.mse_loss, smape, ratio = {'smape':.1}),)

learn.fit(1,.001)
learn.n_beats_trainer.theta_means()<jupyter_output><empty_output><jupyter_text>Example<jupyter_code>from fastseq.data.external import *
items = dummy_data_generator(100, 10, nrows=1000)
data = TSDataLoaders.from_items(items, horizon = 10,lookback = 30, step=5,bs=256)
data.show_batch()
dfs = []
for _ in range(100):
    learn = nbeats_learner(data,layers=[32,], stack_types=("seasonality","trend"), b_loss=.3,
                           loss_func= CombinedLoss(F.mse_loss, smape, ratio = {'smape':.001}),)

    val = learn.validate()[0]

    _df = pd.DataFrame({k:[v.abs().mean().numpy()] for k,v in learn.n_beats_trainer.theta_means(False).items()})
    _df['val'] = val
    dfs.append(_df)
df = pd.concat(dfs)
df
import seaborn as sns
[c for c in list(df.columns) if '_' in c]
df['mean_']= df[[c for c in list(df.columns) if '_' in c]].abs().mean(1)
df['mean_trend1']= df[[c for c in list(df.columns) if 'trend1' in c]].abs().mean(1)
df['mean_seasonality']= df[[c for c in list(df.columns) if 'seasonality' in c]].abs().mean(1)
sns.pairplot(df)
from fastai.callback.all import *
learn.lr_find()
learn.fit(1,1e-12)
learn.fit_one_cycle(10, 2e-3)
learn.recorder.plot_loss()
learn.show_results(max_n=9)
learn.loss_func = CombinedLoss(F.mse_loss, smape, ratio = {'smape':.1})
from fastai.callback.all import *
learn.lr_find()
learn.fit_one_cycle(10, 2e-6)
learn.recorder.plot_loss()
learn.show_results(max_n=9)<jupyter_output><empty_output><jupyter_text>Second Example<jupyter_code>learn = nbeats_learner(data,layers=[200,200,100], theta = .001, b_loss= .1, stack_types=('bias',"trend","seasonality"),
                       loss_func= CombinedLoss(F.mse_loss, smape, ratio = {'smape':.0005}))
from fastai.callback.all import *
learn.lr_find()
learn.fit_one_cycle(10, 2e-3)
learn.recorder.plot_loss()
learn.show_results(max_n=9)
# hide
print(learn.summary())
# hide
from nbdev.export import *
notebook2script()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15060875673144464, 'fraction_numerical': 0.022067431514867713, 'mean_word_length': 3.036144578313253, 'pattern_counts': {'":': 3, '<': 35, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '156652', 'license': 'Apache-2.0', 'n_tokens_mistral': 6430, 'n_tokens_neox': 6193, 'n_words': 1133, 'path': 'nbs/archive/_10_models.nbeats_check_if_theta_mean_diffs_val.ipynb', 'repo_name': 'takotab/fastseq', 'size': '17084'}