conv_cap.py

import mxnet as mx
from mxnet import init
from mxnet import nd
from mxnet.gluon import nn
from mxnet import initializer

def squash(x, axis):
    s_squared_norm = nd.sum(nd.square(x), axis, keepdims=True)
    # if s_squared_norm is really small, we will be in trouble
    # so I removed the s_quare terms
    scale = s_squared_norm / ((1 + s_squared_norm) * nd.sqrt(s_squared_norm + 1e-9))
    return x * scale
    # scale = nd.sqrt(s_squared_norm + 1e-9)
    # return x / scale

class PrimeConvCap(nn.Block):
    def __init__(self, num_cap, num_filter, kernel_size=(3,3),
                  strides=(1,1), padding=(1,1), **kwargs):
        super(PrimeConvCap, self).__init__(**kwargs)
        self.num_cap = num_cap
        self.cap = nn.Conv2D(channels=(num_cap*num_filter), kernel_size=kernel_size,
                             padding=padding, strides=strides)

    def forward(self, x):
        conv_out = nd.expand_dims(self.cap(x), axis=2)
        conv_out = conv_out.reshape((0,self.num_cap,-1,0,0))
        conv_out = squash(conv_out, 2)
        return conv_out


class AdvConvCap(nn.Block):
    def __init__(self, num_cap, num_filter,
                 num_cap_in, num_filter_in,
                 route_num=3, kernel_size=(3,3),
                 strides=(1,1), padding=(1,1),
                 **kwargs):
        super(AdvConvCap, self).__init__(**kwargs)
        self.num_cap = num_cap
        self.num_filter = num_filter
        self.route_num = route_num
        self.num_cap_in = num_cap_in
        # num_filter_in * num_cap_in filters divided in num_cap_in groups
        # with each group output size as num_cap * num_filter
        self.cap = nn.Conv2D(channels=(num_cap * num_filter * num_cap_in), kernel_size=kernel_size, strides=strides,
                             padding=padding, groups= num_cap_in)
        # self.bn = nn.BatchNorm()

    def forward(self, x):
        x_reshape = x.reshape((x.shape[0], -1, x.shape[3], x.shape[4]))
        cap_out = self.cap(x_reshape)
        cap_out = cap_out.reshape((cap_out.shape[0], self.num_cap_in, self.num_cap,
                                   self.num_filter, cap_out.shape[2], cap_out.shape[3]))
        return self.route(cap_out)
        '''
        cap_out = self.cap(x)
        cap_out = cap_out.reshape((x.shape[0],self.num_cap,-1,cap_out.shape[2], cap_out.shape[3]))
        cap_out = squash(cap_out, 2)
        return cap_out
        '''

    def route(self, u):
        b_mat = nd.zeros((u.shape[0], self.num_cap_in, self.num_cap, 1, u.shape[4], u.shape[5]), ctx=u.context)
        for i in range(self.route_num):
            c_mat = nd.softmax(b_mat, axis=2)
            s = nd.sum(u * c_mat, axis=1)
            v = squash(s, 2)
            if i != self.route_num - 1:
                v1 = nd.expand_dims(v, axis=1)
                update_term = nd.sum(u*v1, axis=3, keepdims=True)
                b_mat = b_mat + update_term
        return v

class LengthBlock(nn.Block):
    def __init__(self, **kwargs):
        super(LengthBlock, self).__init__(**kwargs)

    def forward(self, x):
        x = nd.sqrt(nd.sum(nd.square(x), 2))
        return x
'''
class AdvFullyCap(nn.Block):
    def __init__(self, num_cap, num_filter,
                 num_cap_in, num_filter_in,
                 route_num=3, **kwargs):
        self.num_cap = num_cap
        self.num_filter = num_filter
        self.route_num = route_num
        self.num_cap_in = num_cap_in
        # num_filter_in * num_cap_in filters divided in num_cap_in groups
        # with each group output size as num_cap * num_filter
        self.cap = nn.Conv2D(channels=(num_cap * num_filter * num_cap_in), kernel_size=(1,1), groups= num_cap_in)

    def forward(self, x):
        x_reshape = x.reshape((x.shape[0], -1, x.shape[3], x.shape[4]))
        cap_out = self.cap(x_reshape)
        cap_out = cap_out.reshape((cap_out.shape[0], self.num_cap_in, self.num_cap,
                                   self.num_filter, cap_out.shape[2], cap_out.shape[3]))
        return cap_out
'''