(self, dim, processing_steps, num_layers, aggr='add')
| 22 | class Set2SetPool(Pooling): |
| 23 | |
| 24 | def __init__(self, dim, processing_steps, num_layers, aggr='add'): |
| 25 | super(Set2SetPool, self).__init__(aggr=aggr) |
| 26 | self.dim = dim |
| 27 | self.processing_steps = processing_steps |
| 28 | self.num_layers = num_layers |
| 29 | lstm_cells = [tf.nn.rnn_cell.LSTMCell(num_units=dim) |
| 30 | for _ in range(num_layers)] |
| 31 | self.lstm = tf.nn.rnn_cell.MultiRNNCell(lstm_cells) |
| 32 | |
| 33 | def __call__(self, inputs, index, size=None): |
| 34 | size = tf.reduce_max(index) + 1 if size is None else size |
nothing calls this directly
no outgoing calls
no test coverage detected