How to create mini-batches for customized dataset?

[xubingze]

Hi, I am trying to split my own data into mini-batches for better training process. And I am testing the following codes with "Cora" dataset, but the codes do not work. Can you help me figure out the problem? I do not see other people encountering similar issue. Thank you very much!

import argparse
import os.path as osp
import time
import torch
import torch.nn.functional as F
from torch_geometric.loader import LinkNeighborLoader
import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
from torch_geometric.logging import init_wandb, log
from torch_geometric.nn import GATConv
import random
import numpy as np

torch.cuda.empty_cache()
sn = 42
torch.manual_seed(sn)
np.random.seed(sn)
random.seed(sn)
torch.cuda.manual_seed_all(sn)

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='Cora')
parser.add_argument('--hidden_channels', type=int, default=8)
parser.add_argument('--heads', type=int, default=8)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--epochs', type=int, default=200)
parser.add_argument('--wandb', action='store_true', help='Track experiment')
args = parser.parse_args(args=[])

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
init_wandb(name=f'GAT-{args.dataset}', heads=args.heads, epochs=args.epochs,
           hidden_channels=args.hidden_channels, lr=args.lr, device=device)

path = osp.join(osp.dirname(osp.realpath('__file__')), '..', 'data', 'Planetoid')
dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures())
data = dataset[0].to(device)
split = T.RandomLinkSplit(
    num_val=0.05,
    num_test=0.1,
    is_undirected=True,
    add_negative_train_samples=False,
    neg_sampling_ratio=1.0,
)
train_data, val_data, test_data = split(data)

from torch_geometric.loader import DataLoader
batch_size = 128
train_loader = DataLoader(train_data, shuffle=True, batch_size=batch_size)

for batch in train_loader:
  print(batch)

Here is the error:

[rusty1s]

For operating on single graphs, you would want to use NeighborLoader or LinkNeighborLoader instead of DataLoader. In your case, take a look at the LinkNeighborLoader for link prediction tasks.