num_workers = 0
hidden_dims = 16
lr = 3e-5
num_epochs = 5
batches_per_epoch = 20000
g = dataset[0]
// Sampler
batch_sampler = ItemToItemBatchSampler(
g, user_ntype, item_ntype, batch_size)
neighbor_sampler = NeighborSampler(
g, user_ntype, item_ntype, random_walk_length,
random_walk_restart_prob, num_random_walks, num_neighbors,
num_layers)
collator = PinSAGECollator(neighbor_sampler, g, item_ntype, textset)
dataloader = DataLoader(
batch_sampler,
collate_fn=collator.collate_train,
num_workers=num_workers)
dataloader_test = DataLoader(
torch.arange(g.number_of_nodes(item_ntype)),
batch_size=batch_size,
collate_fn=collator.collate_test,
num_workers=num_workers)
dataloader_it = iter(dataloader)
// Model
model = PinSAGEModel(g, item_ntype, textset, hidden_dims, num_layers).to(device)
// Optimizer
opt = torch.optim.Adam(model.parameters(), lr=lr)
model.train()
for batch_id in range(batches_per_epoch):
pos_graph, neg_graph, blocks = next(dataloader_it)
// Copy to GPU
for i in range(len(blocks)):
blocks[i] = blocks[i].to(device)
pos_graph = pos_graph.to(device)
neg_graph = neg_graph.to(device)
loss = model(pos_graph, neg_graph, blocks).mean()
opt.zero_grad()
loss.backward()
opt.step()
print("start training...")
t0 = time.time()
// For each batch of head-tail-negative triplets...
for epoch_id in range(num_epochs):
model.train()
for batch_id in range(batches_per_epoch):
pos_graph, neg_graph, blocks = next(dataloader_it)
// Copy to GPU
for i in range(len(blocks)):
blocks[i] = blocks[i].to(device)
pos_graph = pos_graph.to(device)
neg_graph = neg_graph.to(device)
loss = model(pos_graph, neg_graph, blocks).mean()
opt.zero_grad()
loss.backward()
opt.step()
t1 = time.time()
return (t1 - t0) / num_epochs
After Change
opt = torch.optim.Adam(model.parameters(), lr=lr)
model.train()
for batch_id, (pos_graph, neg_graph, blocks) in enumerate(dataloader):
// Copy to GPU
for i in range(len(blocks)):
blocks[i] = blocks[i].to(device)
pos_graph = pos_graph.to(device)
neg_graph = neg_graph.to(device)
loss = model(pos_graph, neg_graph, blocks).mean()
opt.zero_grad()
loss.backward()
opt.step()
if batch_id >= 3:
break
print("start training...")
t0 = time.time()
// For each batch of head-tail-negative triplets...
for batch_id, (pos_graph, neg_graph, blocks) in enumerate(dataloader):
