How to gather results from different process

[mrbeann]

I try to train a model with ColossalAI and want to calculate some metrics like accuracy. Thus, I record the predictions and truth as follows. But the results show it only contains single batch results. What should I do to get the correct results? (maybe something like torch.distributed.all_gather is needed?)

predictions = []
truths = []
for index, (data_v, labels_v) in enumerate(val_dl):
    data_v, labels_v = data_v.cuda(), labels_v.cuda()
    with torch.no_grad():
        logits = engine(data_v)
    loss = engine.criterion(logits, labels_v)
    running_loss += loss.item()
    final_prob = nn.Softmax(dim=-1)(logits)
    final = final_prob.argmax(dim=-1)
    predictions.append(final)
    truths.append(labels_v)
predictions = np.array(predictions.cpu())
truths = np.array(truths.cpu())
print(predictions.shape, truths.shape, len(val_dl))

[feifeibear]

Thanks for your feedback. I believe you should define your model in a torch.nn.Module class first. The code snippet is not a typical PyTorch Model.

[mrbeann]

Can this API deal with the sync automatically? I do not find any code related to it.

[feifeibear]

What do you mean sync automatically?

[mrbeann]

E.g., in
https://github.com/hpcaitech/ColossalAI/blob/main/examples/images/vit/train.py#L143
if I use a list to record the output.

res = []
for index, (x, y) in tqdm(enumerate(test_dataloader), total=len(test_dataloader), leave=False):
    x, y = x.cuda(), y.cuda()
    output = model(x)
    res.append(output)
    test_loss += F.cross_entropy(output, y, reduction='sum').item()
    pred = output.argmax(dim=1, keepdim=True)
    correct += pred.eq(y.view_as(pred)).sum().item()
    test_sum += y.size(0)
print(res)

Will this work properly?

[feifeibear]

It's supposed to work properly if you update the following line. Since the output may be a shard tensor, we need to convert it to replicate, so its shape is as you expected.

res.append(output.to_replicate())

[mrbeann]

It seems this API uses ColoTensor, so I need another dataloader that also uses ColoTensor?