Map string labels with int label at the end of training


Pytorch doesn’t accept string classes. So, l converted my labels to int and feed my network as follow :

However l need to keep track of my real labels (string names). I need at the end of my learning get the perfomance on each example and map the int label to string label.

How can l do that ?

with open('/home/train.sav', 'rb') as handle:
        train_data = pickle.load(handle)

train_loader =
        train_data, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler)

for i, (input, target) in enumerate(train_loader):

        target = target.cuda(async=True)
        input_var = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)

Thank you

Would a simple dict work?

idx_to_label = {
    0: 'class0',
    1: 'class1',
    2: 'class2'

preds = torch.argmax(torch.randn(10, 2), 1)
for pred in preds:

or what do you want to do with these labels?

1 Like