| 3 | import math |
| 4 | |
| 5 | class DenseGraphConv(Module): |
| 6 | global_count = 0 |
| 7 | def __init__(self, input_channels, output_channels, name=None): |
| 8 | super().__init__() |
| 9 | self.name = (name if name else 'DenseGraph_{}'.format(DenseGraphConv.global_count)) |
| 10 | |
| 11 | DenseGraphConv.global_count+=1 |
| 12 | |
| 13 | bounds = math.sqrt(6.0/(input_channels + output_channels)) |
| 14 | |
| 15 | self.weights_1 = lbann.Weights(initializer = lbann.UniformInitializer(min=-bounds, max=bounds), |
| 16 | name=self.name+'_Weights_1') |
| 17 | self.weights_2 = lbann.Weights(initializer = lbann.UniformInitializer(min=-bounds, max=bounds), |
| 18 | name=self.name+'_Weights_2') |
| 19 | self.W1 = lbann.WeightsLayer(dims = [input_channels, output_channels], |
| 20 | name=self.name+'_param_1', |
| 21 | weights = self.weights_1) |
| 22 | self.W2 = lbann.WeightsLayer(dims = [input_channels, output_channels], |
| 23 | name=self.name+'_param_2', |
| 24 | weights = self.weights_2) |
| 25 | def forward(self, X, A): |
| 26 | messages = lbann.MatMul(X, self.W2, name=self.name+'_w2_mult') |
| 27 | messages = lbann.MatMul(A,messages,name=self.name+'_adj_mult') |
| 28 | |
| 29 | ident = lbann.MatMul(X, self.W1, name=self.name+'_w1_mult') |
| 30 | |
| 31 | out = lbann.Sum(ident, messages, name=self.name+'_sum_id') |
| 32 | |
| 33 | return out |