Optimize construct label function to make it suitable for parallel processing. For loop in my code is creating bottleneck

user_9275386 0
Viewed 34
import torch
import torch.nn as nn

label_tensor = torch.tensor([[0,1,0,0,0,0,0,0,0,0],
                             [0,0,0,0,0,1,0,0,0,0],
                             [0,0,0,0,0,0,0,0,1,0],
                             [0,0,0,0,0,0,1,0,0,0]])

def construct_label(label_tensor):
    batch_list = []
    for i in range(len(label_tensor)):
        empty_tensor_list = []
        for j in label_tensor[i]:
            empty_tensor_list.append(torch.full((28,28),j))
        tensor_label = torch.stack(empty_tensor_list,0)
        batch_list.append(tensor_label)
    batch_tensor_label = torch.stack(batch_list,0)
    print(batch_tensor_label.shape)
    return batch_tensor_label

I want to optimize this function construct_label in more torch way avoiding for loop. Is there a optimized way to do it.

pytorch python
1 Answers

Yes, there is:

import torch
import torch.nn as nn

label_tensor = torch.tensor([[0,1,0,0,0,0,0,0,0,0],
                             [0,0,0,0,0,1,0,0,0,0],
                             [0,0,0,0,0,0,0,0,1,0],
                             [0,0,0,0,0,0,1,0,0,0]])

def construct_label(label_tensor):
    return label_tensor[..., None, None].repeat((1, 1, 28, 28))

You can compare the output of your function with this one using torch.all and you'll see it returns True.

user_42282750
Related