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 hybrid_forward(self, F, X, gram, weight):
self.P = F.batch_dot(F.broadcast_to(weight, shape=(self.gram.shape)),
gram)
if not isinstance(X, mx.symbol.Symbol):
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)
else:
width = X.slice_axis(axis=2, begin=0, end=0)
width = width.ones_like()
width = width.sum()
height = X.slice_axis(axis=3, begin=0, end=0)
height = width.ones_like()
height = width.sum()
print "width", width
print "height", height
arg_shapes, out_shapes, aux_shapes = X.infer_shape_partial(
data=(1, 3, self.width, self.height)) #1 , RGB, Width, Height
return F.batch_dot(
F.SwapAxis(self.P, 1, 2).broadcast_to(
(out_shapes[0][0], self.C, self.C)),
X.reshape((0, 0, out_shapes[0][2] *
out_shapes[0][3]))).reshape(out_shapes[0])
