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

Before Change

    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):

Instances