|
- import torch
- from typing import List
- from .trainprep import PreparedData
- from dataclasses import dataclass
- import random
- from collections import defaultdict
-
-
- @dataclass
- class TrainingBatch(object):
- relation_family_index: int
- relation_type_index: int
- node_type_row: int
- node_type_column: int
- edges: torch.Tensor
-
-
- class FastBatcher(object):
- def __init__(self,
- prep_d: PreparedData,
- batch_size: int) -> None:
-
- if not isinstance(prep_d, PreparedData):
- raise TypeError('prep_d must be an instance of PreparedData')
-
- self.prep_d = prep_d
- self.batch_size = int(batch_size)
-
- self.edges = None
- self.build()
-
- def build(self):
- self.edges = []
- for fam_idx, fam in enumerate(self.prep_d.relation_families):
- edges = []
- targets = []
- edges_back = []
- targets_back = []
- for rel_idx, rel in enumerate(fam.relation_types):
- edges.append(rel.edges_pos.train)
- edges.append(rel.edges_neg.train)
- targets.append(torch.ones(len(rel.edges_pos.train)))
- targets.append(torch.zeros(len(rel.edges_neg.train)))
-
- edges_back.append(rel.edges_back_pos.train)
- edges_back.append(rel.edges_back_neg.train)
- targets_back.apend(torch.zeros(len(rel.edges_back_pos.train)))
- targets_back.apend(torch.zeros(len(rel.edges_back_neg.train)))
-
- edges = torch.cat(edges)
- targets = torch.cat(targets)
- edges_back = torch.cat(edges_back)
- targets_back = torch.cat(targets_back)
-
- order = torch.randperm(len(edges))
- edges = edges[order]
- targets = targets[order]
-
- order_back = torch.randperm(len(edges_back))
- edges_back = edges_back[order_back]
- targets_back = targets_back[order_back]
-
- self.edges.append({'fam_idx': fam_idx, 'rel_idx': rel_idx, 'back': False,
- 'edges': edges, 'targets': targets, 'ofs': 0})
- self.edges.append({'fam_idx': fam_idx, 'rel_idx': rel_idx, 'back': True,
- 'edges': edges_back, 'targets': targets_back, 'ofs': 0})
-
- def __iter__(self):
- while True:
- edges = [ e for e in self.edges \
- if e['ofs'] < len(e['edges']) ]
- # TODO: need to finish this
-
- def __iter_old__(self):
- edge_types = ['edges_pos', 'edges_neg', 'edges_back_pos', 'edges_back_neg']
-
- offsets = {}
- orders = {}
- done = {}
-
- for fam_idx, fam in enumerate(self.prep_d.relation_families):
- for rel_idx, rel in enumerate(fam.relation_types):
- for et in edge_types:
- done[fam_idx, rel_idx, et] = False
-
- while True:
- fam_idx = torch.randint(0, len(self.prep_d.relation_families), (1,)).item()
- fam = self.prep_d.relation_families[fam_idx]
-
- rel_idx = torch.randint(0, len(fam.relation_types), (1,)).item()
- rel = fam.relation_types[rel_idx]
-
- et = random.choice(edge_types)
- edges = getattr(rel, et).train
-
- key = (fam_idx, rel_idx, et)
- if key not in orders:
- orders[key] = torch.randperm(len(edges))
- offsets[key] = 0
-
- ord = orders[key]
- ofs = offsets[key]
-
- nt_row = rel.node_type_row
- nt_col = rel.node_type_column
-
- if 'back' in et:
- nt_row, nt_col = nt_col, nt_row
-
- if ofs < len(edges):
- offsets[key] += self.batch_size
- ord = ord[ofs:ofs+self.batch_size]
- edges = edges[ord]
- yield TrainingBatch(fam_idx, rel_idx, nt_row, nt_column, edges)
- else:
- done[key] = True
-
-
-
-
- for fam in self.prep_d.relation_families:
- edges = []
- for rel in fam.relation_types:
- edges.append(rel.edges_pos.train)
- edges.append(rel.edges_back_pos.train)
- edges.append(rel.edges_neg.train)
- edges.append(rel.edges_back_neg.train)
- edges = torch.cat(e)
-
-
-
- class FastDecLayer(torch.nn.Module):
- def __init__(self, **kwargs):
- super().__init__(**kwargs)
-
- def forward(self,
- last_layer_repr: List[torch.Tensor],
- training_batch: TrainingBatch):
|
