22ccf4365af620d10387b207aa103287c34d9247,benchmarks/benchmarks/model_speed/bench_pinsage.py,,track_time,#Any#,363

Before Change

    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)

After Change

    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):
        // 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 >= 10:  // time 10 loops
            break

    t1 = time.time()

    return (t1 - t0) / (batch_id + 1)

Instances