Lightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Go, and more
class Inspiration(HybridBlock):
""" Inspiration Layer (from MSG-Net paper)
tuning the featuremap with target Gram Matrix
ref https://arxiv.org/abs/1703.06953
"""
def __init__(self, C, B=1):
super(Inspiration, self).__init__()
# B is equal to 1 or input mini_batch
self.C = C
self.weight = self.params.get('weight', shape=(1,C,C),
init=mx.initializer.Uniform(),
allow_deferred_init=True)
self.gram = self.params.get('gram', shape=(B,C,C),
init=mx.initializer.Uniform(),
allow_deferred_init=True,
lr_mult=0)
def setTarget(self, target):
self.gram.set_data(target)
def forward(self, X):
# input X is a 3D feature map
self.P = F.batch_dot(F.broadcast_to(self.weight.data(), shape=(self.gram.shape)), self.gram.data())
return F.batch_dot(F.SwapAxis(self.P,1,2).broadcast_to((X.shape[0], self.C, self.C)), X.reshape((0,0,X.shape[2]*X.shape[3]))).reshape(X.shape)
def __repr__(self):
return self.__class__.__name__ + '(' \
+ 'N x ' + str(self.C) + ')'
How do you implement the hybrid forward function for this block, given X.shape is not a method of a Symbol and X.infer_shape(Y), needs an argument of Y as a shape and returns args_shapes,output_shapes,input_shapes # B is equal to 1 or input mini_batch
self.C = C
self.B = B
self.weight = self.collect_params().get('weight', shape=(1,self.C,self.C),
init=mx.initializer.Uniform(),
allow_deferred_init=True)
self.gram = self.collect_params().get('gram', shape=(self.B,self.C,self.C),
init=mx.initializer.Uniform(),
allow_deferred_init=True,
lr_mult=0)
self.weight.initialize(ctx=ctx)
self.gram.initialize(ctx=ctx)
def setTarget(self, target):
self.gram.set_data(target)
def hybrid_forward(self, F, X, gram, weight):
P = F.batch_dot(F.broadcast_to(weight, shape=(self.gram.shape)), gram)
if not isinstance(X,symbol.Symbol):
return F.batch_dot(F.SwapAxis(P,1,2).broadcast_to((X.shape[0], self.C, self.C)), X.reshape((0,0,X.shape[2]*X.shape[3]))).reshape(X.shape)
else:
#print "Hooppla", interals
#for i in dir(interals):
# print "kk:", i
in_shapes,out_shapes,arg_shapes= X.infer_shape(self.gram.shape)
#print out_shapes
#raise Exception
#arg_shapes, out_shapes, aux_shapes = interals.infer_shape(self.gram.shape)
#print "A", arg_shapes, "O", out_shapes, "AU", aux_shapes
return F.batch_dot(F.SwapAxis(P,1,2).broadcast_to((in_shapes[0], self.C, self.C)), X.reshape((0,0,in_shapes[2]*in_shapes[3]))).reshape(in_shapes)
def __repr__(self):
return self.__class__.__name__ + '(' \
+ 'N x ' + str(self.C) + ')'`