|
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576 |
- #
- # 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 = {}
-
- for (node_type_row, node_type_column), rels in self.data.relation_types.items():
- 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, k) for k in range(dec.num_relation_types) ]
- res[node_type_row, node_type_column] = pred_adj_matrices
- return res
|