masked_pool_lasagne.py

"""symbolic helper functions"""

import theano
import theano.tensor as T
import lasagne
from lasagne.layers import *
import numpy as np

def get_mask(tokens,pad_ix=-1):
    assert tokens.ndim==2
    return T.eq(T.cumsum(T.eq(tokens,pad_ix),axis=1),0)

def mellowmax(x,w=5,axis=2,mask=None):
    exp = T.exp(w*x)
    mean_exp = exp.mean(axis) if mask is None else (exp*mask).sum(axis)/mask.sum(axis)
    return T.log(mean_exp)/w
                         
class MaskedGlobalPoolLayer(MergeLayer):
    def __init__(self, 
                 incoming, 
                 pool_function=T.mean,
                 mask_input=None,**kwargs):
        incomings = [incoming]
        if mask_input is not None:
            incomings.append(mask_input)
        MergeLayer.__init__(self, incomings , **kwargs)
        self.pool_function = pool_function
    def get_output_shape_for(self, input_shapes):
        return input_shapes[0][:2]

    def get_output_for(self, inputs, **kwargs):
        x = inputs[0].flatten(3)
        if len(inputs)==2:
            mask = inputs[1]
            mask = mask[:,None,-x.shape[2]:]
            return self.pool_function(x,mask=mask,axis=2)
        else:
            return self.pool_function(x,axis=2)
        

class MaskedKMaxPoolLayer(MergeLayer):
    def __init__(self, 
                 incoming, 
                 k=3,
                 lower_bound=-np.inf,
                 mask_input=None,**kwargs):
        incomings = [incoming]
        self.k=k
        self.lower_bound=lower_bound
        if mask_input is not None:
            incomings.append(mask_input)
        MergeLayer.__init__(self, incomings , **kwargs)
        
    def get_output_shape_for(self, input_shapes):
        
        batch,units = input_shapes[0][:2]
        return (batch,units,self.k)

    def get_output_for(self, inputs, **kwargs):
        data = inputs[0].flatten(3)
        if len(inputs)==2:
            mask = inputs[1]
            mask = mask[:,None,-data.shape[2]:]
            data =  T.switch(T.eq(mask, 0), self.lower_bound, data)
            
        k_max_ix = T.sort(T.argsort(data, axis=2)[:,:,-self.k:],axis=2)

        return data[T.arange(data.shape[0])[:,None,None],T.arange(data.shape[1])[None,:,None],k_max_ix]