I want to implement the mini-batch training of “node v.s. node” contrastive learning between two views of the same graph ? How easy should it be to implement?
The code I’m using now is as follows, and it’s slow. I hope it can be improved.
indices = list(range(graph.num_nodes()))
np.random.shuffle(indices)
shuffle_indices = torch.tensor(indices).to(graph.device)
sampler = dgl.dataloading.MultiLayerFullNeighborSampler(num_layers=2)
dataloader1 = dgl.dataloading.DataLoader(g1, shuffle_indices, sampler, batch_size=batch_size, drop_last=False, shuffle=False, num_workers=4)
dataloader2 = dgl.dataloading.DataLoader(g2, shuffle_indices, sampler, batch_size=batch_size,drop_last=False, shuffle=False, num_workers=4)
for x, y in zip(dataloader1, dataloader2):
assert torch.equal(x[1], y[1])
blocks1, blocks2 = x[-1], y[-1]
blocks1 = [b.to(device) for b in blocks1]
blocks2 = [b.to(device) for b in blocks2]
z1 = model.embedding(blocks1, blocks1[0].srcdata['feat'])
z2 = model.embedding(blocks2, blocks2[0].srcdata['feat'])
assert torch.isnan(z1).sum()==0
assert torch.isnan(z2).sum()==0
loss = model.get_loss([z1, z2])
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_epoch += loss.item()