Fit a set of independent HeteroData to a heterogeneous model for classification.

[LynchOrange]

from torch.nn import Linear
import torch
from torch_geometric.loader import DataListLoader
from torch_geometric.nn import GCNConv, SAGEConv, GATConv, DataParallel,HeteroConv, GlobalAttention
from htg_data import HTG_data
import torch.nn.functional as F

class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_layers):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                ('token', 'next_token', 'token'): GCNConv(-1, hidden_channels),
                ('token', 'token_sink', 'sink'): SAGEConv((-1, -1), hidden_channels),
                ('token', 'belongs_to', 'property'): GATConv((-1, -1), hidden_channels),
                ('property', 'property_sink', 'sink'): SAGEConv((-1, -1), hidden_channels),
                ('property', 'next_property', 'property'): GATConv((-1, -1), hidden_channels),
            }, aggr='sum')
            self.convs.append(conv)
        self.pooling_gate_nn = Linear(hidden_channels, 1)
        self.pooling = GlobalAttention(self.pooling_gate_nn)
        self.lin = Linear(hidden_channels, out_channels)

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        self.pooling.reset_parameters()
        self.lin.reset_parameters()

    def forward(self, data):
        x_dict, edge_index_dict, batch = data.x_dict, data.edge_index_dict, data.batch
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        x = torch.cat((x_dict['sink'],x_dict['property'],x_dict['token']), 0)
        x = self.pooling(x, batch)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin(x)
        return x

dataset = HTG_data(root = '/home/ylzqn/HTG_Data/HTG_CAG')

print(dataset[0])

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

train_loader = DataListLoader(dataset = dataset[0:1000], batch_size = 128, shuffle = True)

model = HeteroGNN(hidden_channels=64, out_channels=2, num_layers=1)

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = DataParallel(model)

model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.00005)
criterion = torch.nn.MultiLabelSoftMarginLoss()


def train():
    model.train()

    for data_list in train_loader:  
        out = model(data_list)
        y = torch.cat([torch.unsqueeze(data.y, 0).float() for data in data_list], 0).to(device)
        loss = criterion(out, y)
        loss.backward()  
        optimizer.step()
        optimizer.zero_grad() 
    print("  Loss: ", loss.data.mean())
    print("Train Done")

for epoch in range(1, 40):
    train()

---

HeteroData(
  y=[2],
  token={ x=[67, 768] },
  property={ x=[169, 49] },
  sink={ x=[1, 768] },
  (token, next_token, token)={ edge_index=[2, 66] },
  (token, token_sink, sink)={ edge_index=[2, 67] },
  (token, belongs_to, property)={ edge_index=[2, 67] },
  (property, property_sink, sink)={ edge_index=[2, 3] },
  (property, next_property, property)={ edge_index=[2, 166] }
)

Traceback (most recent call last):
  File "/home/ylzqn/Jupyter Notebook/pkgcode2vec/htg_model.py", line 96, in <module>
    train()
  File "/home/ylzqn/Jupyter Notebook/pkgcode2vec/htg_model.py", line 86, in train
    out = model(data_list)  # Perform a single forward pass.
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/nn/data_parallel.py", line 67, in forward
    inputs = self.scatter(data_list, self.device_ids)
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/nn/data_parallel.py", line 92, in scatter
    for i in range(len(split) - 1)
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/nn/data_parallel.py", line 92, in <listcomp>
    for i in range(len(split) - 1)
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/data/batch.py", line 69, in from_data_list
    exclude_keys=exclude_keys,
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/data/collate.py", line 77, in collate
    increment)
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/data/collate.py", line 144, in _collate
    values = [value + inc for value, inc in zip(values, incs)]
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch_geometric/data/collate.py", line 144, in <listcomp>
    values = [value + inc for value, inc in zip(values, incs)]
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and CPU!

---

Is that possible to use a set of HeteroData as the input to a heterogeneous model for classifying the class of each HeteroData?
The DataListLoader for object HeteroData is different from it for object Data. I can't directly extract the attributes from HeteroDataBatch. For example:
If the object of DataListLoader is Data, I can get node features by “x = batch.x”, but if the object of DataListLoader is HeteroData, the code "x_dict = batch.x_dict" doesn't work. Is there an example of applying a heterogeneous model on multiple graphs?

[rusty1s]

DataListLoader currently does not work with HeteroData objects, I am sorry. For multi-GPU training, it is recommended to use DistributedDataParallel instead, see here.

[LynchOrange]

import os
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn import Linear
import torch
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data.distributed import DistributedSampler
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, SAGEConv, GATConv, HeteroConv, GlobalAttention
from htg_data import HTG_data
import torch.nn.functional as F

class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_layers=2, dropout=0.5):
        super().__init__()
        self.dropout = dropout
        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                ('token', 'next_token', 'token'): GCNConv(-1, hidden_channels),
                ('token', 'token_sink', 'sink'): SAGEConv((-1, -1), hidden_channels),
                ('token', 'belongs_to', 'property'): GATConv((-1, -1), hidden_channels),
                ('property', 'property_sink', 'sink'): SAGEConv((-1, -1), hidden_channels),
                ('property', 'next_property', 'property'): GATConv((-1, -1), hidden_channels),
            }, aggr='sum')
            self.convs.append(conv)
        self.pooling_gate_nn = Linear(hidden_channels, 1)
        self.pooling = GlobalAttention(self.pooling_gate_nn)
        self.lin = Linear(hidden_channels, out_channels)

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        self.pooling.reset_parameters()
        self.lin.reset_parameters()

    def forward(self, x_dict, edge_index_dict, batch):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        x = torch.cat((x_dict['sink'],x_dict['property'],x_dict['token']), 0)
        x = self.pooling(x, batch)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.lin(x)
        return x

def run(rank, world_size: int, root: str):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

    dataset = HTG_data(root = root)

    print(dataset[0])

    train_sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    train_loader = DataLoader(dataset, batch_size=128, sampler=train_sampler)


    torch.manual_seed(12345)
    model = HeteroGNN(hidden_channels=64, out_channels=2, num_layers=2).to(rank)
    model = DistributedDataParallel(model, device_ids=[rank])

    print(model)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
    criterion = torch.nn.MultiLabelSoftMarginLoss()

    for epoch in range(1, 51):
        model.train()

        total_loss = 0
        for data in train_loader:
            data = data.to(rank)
            optimizer.zero_grad()
            logits = model(data.x_dict, data.edge_index_dict, data.batch)
            loss = criterion(logits, data.y.to(torch.float))
            loss.backward()
            optimizer.step()
            total_loss += float(loss) * logits.size(0)
        loss = total_loss / len(train_loader.dataset)

        dist.barrier()

        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, ')

    dist.destroy_process_group()

if __name__ == '__main__':

    root = '/home/ylzqn/HTG_Data/HTG_CAG'
    world_size = torch.cuda.device_count()
    print('Let\'s use', world_size, 'GPUs!')
    args = (world_size, root)
    mp.spawn(run, args=args, nprocs=world_size, join=True)

---

Traceback (most recent call last):
  File "/home/ylzqn/Jupyter Notebook/pkgcode2vec/htg_model.py", line 127, in <module>
    mp.spawn(run, args=args, nprocs=world_size, join=True)
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
    while not context.join():
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 150, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
    fn(i, *args)
  File "/home/ylzqn/Jupyter Notebook/pkgcode2vec/htg_model.py", line 71, in run
    model = DistributedDataParallel(model, device_ids=[rank])
  File "/home/ylzqn/.conda/envs/lynch_pytorch/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 435, in __init__
    'Modules with uninitialized parameters can\'t be used with `DistributedDataParallel`. '
RuntimeError: Modules with uninitialized parameters can't be used with `DistributedDataParallel`. Run a dummy forward pass to correctly initialize the modules

---

Thank you for your reply. When I followed the example you provided, I got a new error. I have tried to initialize all parameters, but it still doesn't work.

[rusty1s]

Since you are using lazy initialization to create your model (the -1 values in the GNN layers), the parameters need to be first materialized before inputting the model into distributed data parallel mode. You can either run a forward pass, or initialize them directly by providing the correct feature dimension instead of -1.