Class Collater

model-experiments/gnn-based-experiments/src/data/dataloader.py:10–56

# from src.data.batch import Batch

class Collater(object):
    def __init__(self, follow_batch, ndevices):
        self.follow_batch = follow_batch
        self.ndevices = ndevices

    def collate(self, batch):
        elem = batch[0]
        if isinstance(elem, Data):
            data = batch
            count = torch.tensor([data.num_nodes for data in data])
            cumsum = count.cumsum(0)
            cumsum = torch.cat([cumsum.new_zeros(1), cumsum], dim=0)
            device_id = self.ndevices * cumsum.to(torch.float) / cumsum[-1].item()
            device_id = (device_id[:-1] + device_id[1:]) / 2.0
            device_id = device_id.to(torch.long)  # round.
            split = device_id.bincount().cumsum(
                0)  # Count the frequency of each value in an array of non-negative ints.
            split = torch.cat([split.new_zeros(1), split], dim=0)
            split = torch.unique(split, sorted=True)
            split = split.tolist()

            graphs = []
            for i in range(len(split) - 1):
                data1 = data[split[i]:split[i + 1]]
                graph = Batch.from_data_list(data1, self.follow_batch)
                graphs += [graph]

            return graphs  #Batch.from_data_list(batch, self.follow_batch)
        elif isinstance(elem, torch.Tensor):
            return default_collate(batch)
        elif isinstance(elem, float):
            return torch.tensor(batch, dtype=torch.float)
        elif isinstance(elem, int_classes):
            return torch.tensor(batch)
        elif isinstance(elem, string_classes):
            return batch
        elif isinstance(elem, container_abcs.Mapping):
            return {key: self.collate([d[key] for d in batch]) for key in elem}
        elif isinstance(elem, tuple) and hasattr(elem, '_fields'):
            return type(elem)(*(self.collate(s) for s in zip(*batch)))
        elif isinstance(elem, container_abcs.Sequence):
            return [self.collate(s) for s in zip(*batch)]

        raise TypeError('DataLoader found invalid type: {}'.format(type(elem)))

    def __call__(self, batch):
        return self.collate(batch)


class DataLoader(torch.utils.data.DataLoader):