graphlearn_torch/python/loader/link_loader.py

# Copyright 2022 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================


from typing import Tuple, Union, Optional, Literal

import torch

from .transform import to_data, to_hetero_data

from ..utils import convert_to_tensor
from ..data import Dataset
from ..sampler import (
  BaseSampler,
  EdgeSamplerInput,
  NegativeSampling,
  SamplerOutput,
  HeteroSamplerOutput
)
from ..typing import InputEdges, reverse_edge_type


class LinkLoader(object):
  r"""A data loader that performs mini-batch sampling from link information,
  using a generic :class:`~graphlearn_torch.sampler.BaseSampler`
  implementation that defines a
  :meth:`~graphlearn_torch.sampler.BaseSampler.sample_from_edges` function and
  is supported on the provided input :obj:`data` object.

    .. note::
      Negative sampling for triplet case is currently implemented in an
      approximate way, *i.e.* negative edges may contain false negatives.

    Args:
      data (Dataset): The `graphlearn_torch.data.Dataset` object.
      link_sampler (graphlearn_torch.sampler.BaseSampler): The sampler
        implementation to be used with this loader.
        Needs to implement
        :meth:`~graphlearn_torch.sampler.BaseSampler.sample_from_edges`.
        The sampler implementation must be compatible with the input
        :obj:`data` object.
      edge_label_index (Tensor or EdgeType or Tuple[EdgeType, Tensor]):
        The edge indices, holding source and destination nodes to start
        sampling from.
        If set to :obj:`None`, all edges will be considered.
        In heterogeneous graphs, needs to be passed as a tuple that holds
        the edge type and corresponding edge indices.
        (default: :obj:`None`)
      edge_label (Tensor, optional): The labels of edge indices from which to
        start sampling from. Must be the same length as
        the :obj:`edge_label_index`. (default: :obj:`None`)
      neg_sampling (NegativeSampling, optional): The negative sampling
        configuration.
        For negative sampling mode :obj:`"binary"`, samples can be accessed
        via the attributes :obj:`edge_label_index` and :obj:`edge_label` in
        the respective edge type of the returned mini-batch.
        In case :obj:`edge_label` does not exist, it will be automatically
        created and represents a binary classification task (:obj:`0` =
        negative edge, :obj:`1` = positive edge).
        In case :obj:`edge_label` does exist, it has to be a categorical
        label from :obj:`0` to :obj:`num_classes - 1`.
        After negative sampling, label :obj:`0` represents negative edges,
        and labels :obj:`1` to :obj:`num_classes` represent the labels of
        positive edges.
        Note that returned labels are of type :obj:`torch.float` for binary
        classification (to facilitate the ease-of-use of
        :meth:`F.binary_cross_entropy`) and of type
        :obj:`torch.long` for multi-class classification (to facilitate the
        ease-of-use of :meth:`F.cross_entropy`).
        For negative sampling mode :obj:`"triplet"`, samples can be
        accessed via the attributes :obj:`src_index`, :obj:`dst_pos_index`
        and :obj:`dst_neg_index` in the respective node types of the
        returned mini-batch.
        :obj:`edge_label` needs to be :obj:`None` for :obj:`"triplet"`
        negative sampling mode.
        If set to :obj:`None`, no negative sampling strategy is applied.
        (default: :obj:`None`)
      device (torch.device, optional): The device to put the data on.
        If set to :obj:`None`, the CPU is used.
      edge_dir (str:["in", "out"]): The edge direction for sampling.
        Can be either :str:`"out"` or :str:`"in"`.
        (default: :str:`"out"`)
      **kwargs (optional): Additional arguments of
        :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
        :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.

  """
  def __init__(
    self,
    data: Dataset,
    link_sampler: BaseSampler,
    edge_label_index: InputEdges = None,
    edge_label: Optional[torch.Tensor] = None,
    neg_sampling: Optional[NegativeSampling] = None,
    device: torch.device = torch.device('cuda:0'),
    edge_dir: Literal['out', 'in'] = 'out',
    **kwargs,
  ):
    # Get edge type (or `None` for homogeneous graphs):
    input_type, edge_label_index = get_edge_label_index(
        data, edge_label_index)
    self.data = data
    self.link_sampler = link_sampler
    self.neg_sampling = NegativeSampling.cast(neg_sampling)
    self.device = device
    self.edge_dir = edge_dir

    if (self.neg_sampling is not None and self.neg_sampling.is_binary()
        and edge_label is not None and edge_label.min() == 0):
      # Increment labels such that `zero` now denotes "negative".
      edge_label = edge_label + 1

    if (self.neg_sampling is not None and self.neg_sampling.is_triplet()
            and edge_label is not None):
      raise ValueError("'edge_label' needs to be undefined for "
                        "'triplet'-based negative sampling. Please use "
                        "`src_index`, `dst_pos_index` and "
                        "`neg_pos_index` of the returned mini-batch "
                        "instead to differentiate between positive and "
                        "negative samples.")

    self.input_data = EdgeSamplerInput(
      row=edge_label_index[0].clone(),
      col=edge_label_index[1].clone(),
      label=edge_label,
      input_type=input_type,
      neg_sampling=self.neg_sampling,
    )

    input_index = range(len(edge_label_index[0]))
    self._seed_loader = torch.utils.data.DataLoader(input_index, **kwargs)

  def __iter__(self):
    self._seeds_iter = iter(self._seed_loader)
    return self

  def __next__(self):
    seeds = self._seeds_iter._next_data().to(self.device)
    # Currently, we support the out-edge sampling manner, so we reverse the
    # direction of src and dst for the output so that features of the sampled
    # nodes during training can be aggregated from k-hop to (k-1)-hop nodes.
    sampler_out = self.link_sampler.sample_from_edges(self.input_data[seeds])
    result = self._collate_fn(sampler_out)

    return result

  def _collate_fn(self, sampler_out: Union[SamplerOutput, HeteroSamplerOutput]):
    r"""format sampler output to Data/HeteroData
      For the out-edge sampling scheme (i.e. the direction of edges in
      the output is inverse to the original graph), we put the reversed
      edge_label_index into the (dst, rev_to, src) subgraph for
      HeteroSamplerOutput and (dst, to, src) for SamplerOutput.
      However, for the in-edge sampling scheme (i.e. the direction of edges 
      in the output is the same as the original graph), we do not need to
      reverse the edge type of the sampler_out.
    """
    if isinstance(sampler_out, SamplerOutput):
      x = self.data.node_features[sampler_out.node]
      if self.data.edge_features is not None and sampler_out.edge is not None:
        edge_attr = self.data.edge_features[sampler_out.edge]
      else:
        edge_attr = None
      res_data = to_data(sampler_out,
                         node_feats=x,
                         edge_feats=edge_attr,
                        )
    else: # hetero
      x_dict = {}
      x_dict = {ntype : self.data.get_node_feature(ntype)[ids.to(torch.int64)] for ntype, ids in sampler_out.node.items()}
      edge_attr_dict = {}
      if sampler_out.edge is not None:
        for etype, eids in sampler_out.edge.items():
          if self.edge_dir == 'out':
            efeat = self.data.get_edge_feature(reverse_edge_type(etype))
          elif self.edge_dir == 'in':
            efeat = self.data.get_edge_feature(etype)
          if efeat is not None:
            edge_attr_dict[etype] = efeat[eids.to(torch.int64)]

      res_data = to_hetero_data(sampler_out,
                                node_feat_dict=x_dict,
                                edge_feat_dict=edge_attr_dict,
                                edge_dir=self.edge_dir,
                               )
    return res_data

  def __repr__(self) -> str:
    return f'{self.__class__.__name__}()'


def get_edge_label_index(
  data: Dataset,
  edge_label_index: InputEdges
) -> Tuple[Optional[str], torch.Tensor]:
  edge_type = None
  # # Need the edge index in COO for LinkNeighborLoader:
  def _get_edge_index(edge_type):
    row, col, _, _ = data.get_graph(edge_type).topo.to_coo()
    return (row, col)

  if not isinstance(edge_label_index, Tuple):
    if edge_label_index is None:
      return None, _get_edge_index(edge_type)
    return None, convert_to_tensor(edge_label_index)

  if isinstance(edge_label_index[0], str):
    edge_type = edge_label_index

    return edge_type, _get_edge_index(edge_type)

  assert len(edge_label_index) == 2
  edge_type, edge_label_index = convert_to_tensor(edge_label_index)

  if edge_label_index is None:
    row, col, _, _ = data.get_graph(edge_type).topo.to_coo()
    return edge_type, (row, col)

  return edge_type, edge_label_index