Initialize GatedGraph layer Args: input_channels (int): The size of the input node features output_channels (int): The output size of the node features num_nodes (int): Number of vertices in the graph num_layers (int): Number of passes through
(self,
input_channels,
output_channels,
num_nodes,
num_layers = 1,
name = None)
| 16 | """ |
| 17 | global_count = 0 |
| 18 | def __init__(self, |
| 19 | input_channels, |
| 20 | output_channels, |
| 21 | num_nodes, |
| 22 | num_layers = 1, |
| 23 | name = None): |
| 24 | """Initialize GatedGraph layer |
| 25 | Args: |
| 26 | input_channels (int): The size of the input node features |
| 27 | output_channels (int): The output size of the node features |
| 28 | num_nodes (int): Number of vertices in the graph |
| 29 | num_layers (int): Number of passes through the GRU (default: 1) |
| 30 | name (str): Name of the layers and prefix to use for the layers. |
| 31 | data_layout (str): Data layout (default: data parallel) |
| 32 | """ |
| 33 | super().__init__() |
| 34 | |
| 35 | ## Add Name for the components for the layer |
| 36 | GatedGraphConv.global_count +=1 |
| 37 | self.name = (name |
| 38 | if name |
| 39 | else 'GatedGraphConv_{}'.format(GatedGraphConv.global_count)) |
| 40 | |
| 41 | |
| 42 | ## Add variables |
| 43 | self.output_channel_size = output_channels |
| 44 | self.input_channel_size = input_channels |
| 45 | self.num_nodes = num_nodes |
| 46 | |
| 47 | self.rnn = lbann.modules.ChannelwiseGRU(num_nodes, output_channels) |
| 48 | |
| 49 | self.num_layers = num_layers |
| 50 | self.nns = [] |
| 51 | |
| 52 | for i in range(num_layers): |
| 53 | |
| 54 | weights = lbann.Weights(initializer = lbann.UniformInitializer(min =-1/(math.sqrt(output_channels)), |
| 55 | max = 1/(math.sqrt(output_channels)))) |
| 56 | nn = \ |
| 57 | ChannelwiseFullyConnectedModule(self.output_channel_size, |
| 58 | bias=False, |
| 59 | weights=[weights], |
| 60 | name=f"{self.name}_nn_{i}") |
| 61 | self.nns.append(nn) |
| 62 | |
| 63 | |
| 64 | def forward(self, node_feature_mat, source_indices, target_indices): |
nothing calls this directly
no test coverage detected