|
- from .data import Data
- from .model import TrainingBatch
- import torch
-
-
- def _shuffle(x: torch.Tensor) -> torch.Tensor:
- order = torch.randperm(len(x))
- return x[order]
-
-
- class Batcher(object):
- def __init__(self, data: Data, batch_size: int=512,
- shuffle: bool=True) -> None:
-
- if not isinstance(data, Data):
- raise TypeError('data must be an instance of Data')
-
- self.data = data
- self.batch_size = int(batch_size)
- self.shuffle = bool(shuffle)
-
- def __iter__(self) -> TrainingBatch:
- edge_types = list(self.data.edge_types.values())
-
- edge_lists = [ [ adj_mat.indices().transpose(0, 1) \
- for adj_mat in et.adjacency_matrices ] \
- for et in edge_types ]
-
- if self.shuffle:
- edge_lists = [ [ _shuffle(lst) for lst in edge_lst ] \
- for edge_lst in edge_lists ]
-
- offsets = [ [ 0 ] * len(et.adjacency_matrices) \
- for et in edge_types ]
-
- while True:
- candidates = [ edge_idx for edge_idx, edge_ofs in enumerate(offsets) \
- if len([ rel_idx for rel_idx, rel_ofs in enumerate(edge_ofs) \
- if rel_ofs < len(edge_lists[edge_idx][rel_idx]) ]) > 0 ]
- if len(candidates) == 0:
- break
-
- edge_idx = torch.randint(0, len(candidates), (1,)).item()
- edge_idx = candidates[edge_idx]
- candidates = [ rel_idx \
- for rel_idx, rel_ofs in enumerate(offsets[edge_idx]) \
- if rel_ofs < len(edge_lists[edge_idx][rel_idx]) ]
-
- rel_idx = torch.randint(0, len(candidates), (1,)).item()
- rel_idx = candidates[rel_idx]
-
- lst = edge_lists[edge_idx][rel_idx]
- et = edge_types[edge_idx]
- ofs = offsets[edge_idx][rel_idx]
- lst = lst[ofs:ofs+self.batch_size]
- offsets[edge_idx][rel_idx] += self.batch_size
-
- b = TrainingBatch(et.vertex_type_row, et.vertex_type_column,
- rel_idx, lst, torch.full((len(lst),), self.data.target_value,
- dtype=torch.float32))
-
- yield b
|