implement scatter general readout
@@ -38,3 +38,24 @@ class GeneralReadout(chainer.Link):
else:
raise ValueError('mode {} is not supported'.format(self.mode))
return y
+
+
+class ScatterGeneralReadout(chainer.Link):
+ def __init__(self, mode='sum', activation=None, **kwargs):
+ super(ScatterGeneralReadout, self).__init__()
+ self.mode = mode
+ self.activation = activation
+
+ def __call__(self, h, batch, axis=1, **kwargs):
+ if self.activation is not None:
+ h = self.activation(h)
+ else:
+ h = h
+
+ if self.mode == 'sum':
+ y = self.xp.zeros((batch[-1] + 1, h.shape[1]),
+ dtype=self.xp.float32)
+ y = functions.scatter_add(y, batch, h)
+ else:
+ raise ValueError('mode {} is not supported'.format(self.mode))
+ return y
