|
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586 |
- #
- # Copyright (C) Stanislaw Adaszewski, 2020
- # License: GPLv3
- #
-
-
- import torch
- from .data import Data
- from .trainprep import PreparedData, \
- TrainValTest
- from typing import Type, \
- List, \
- Callable, \
- Union, \
- Dict, \
- Tuple
- from .decode import DEDICOMDecoder
-
-
- class DecodeLayer(torch.nn.Module):
- def __init__(self,
- input_dim: List[int],
- data: Union[Data, PreparedData],
- keep_prob: float = 1.,
- decoder_class: Union[Type, Dict[Tuple[int, int], Type]] = DEDICOMDecoder,
- activation: Callable[[torch.Tensor], torch.Tensor] = torch.sigmoid,
- **kwargs) -> None:
-
- super().__init__(**kwargs)
-
- assert all([ a == input_dim[0] \
- for a in input_dim ])
-
- self.input_dim = input_dim
- self.output_dim = 1
- self.data = data
- self.keep_prob = keep_prob
- self.decoder_class = decoder_class
- self.activation = activation
-
- self.decoders = None
- self.build()
-
- def build(self) -> None:
- self.decoders = {}
-
- n = len(self.data.node_types)
- relation_types = self.data.relation_types
- for node_type_row in range(n):
- if node_type_row not in relation_types:
- continue
-
- for node_type_column in range(n):
- if node_type_column not in relation_types[node_type_row]:
- continue
-
- rels = relation_types[node_type_row][node_type_column]
- if len(rels) == 0:
- continue
-
- if isinstance(self.decoder_class, dict):
- if (node_type_row, node_type_column) in self.decoder_class:
- decoder_class = self.decoder_class[node_type_row, node_type_column]
- elif (node_type_column, node_type_row) in self.decoder_class:
- decoder_class = self.decoder_class[node_type_column, node_type_row]
- else:
- raise KeyError('Decoder not specified for edge type: %s -- %s' % (
- self.data.node_types[node_type_row].name,
- self.data.node_types[node_type_column].name))
- else:
- decoder_class = self.decoder_class
-
- self.decoders[node_type_row, node_type_column] = \
- decoder_class(self.input_dim[node_type_row],
- num_relation_types = len(rels),
- keep_prob = self.keep_prob,
- activation = self.activation)
-
- def forward(self, last_layer_repr: List[torch.Tensor]) -> Dict[Tuple[int, int], List[torch.Tensor]]:
- res = {}
- for (node_type_row, node_type_column), dec in self.decoders.items():
- inputs_row = last_layer_repr[node_type_row]
- inputs_column = last_layer_repr[node_type_column]
- pred_adj_matrices = dec(inputs_row, inputs_column)
- res[node_type_row, node_type_column] = pred_adj_matrices
- return res
|