Tensor Manipulation for Padding

I’m working with graph neural networks and it’s important for my work to perform some padding. Here is the problem -

I have 3 graphs [g1, g2, g3] and each with 1,2,3 nodes respectively. I want to aggregate these nodes such that they are padded with a fixed size. Here is the example

all_nodes_feature = torch.tensor[ 
[1, 1],
[2, 2],
[2, 2],
[3, 3],
[3, 3],
[3, 3] ]

number_of_nodes_in_each_graph = torch.tensor([1, 2, 3])

fixed_length = 3

resulting_output = [ 
[ [1,1], [0, 0], [0,0] ],
[ [2,2], [2, 2], [0,0] ],
[ [3,3], [3, 3], [3,3] ]
]

Hi,

you could firstly split the all_nodes_feature tensor to then pad them as sequences. Be aware however that in this case number_of_nodes_in_each_graph should be a normal list of integers.
So you would have:

all_nodes_feature = torch.tensor([[1, 1],
                                 [2, 2],
                                 [2, 2],
                                 [3, 3],
                                 [3, 3],
                                 [3, 3]])

number_of_nodes_in_each_graph = [1, 2, 3]

all_nodes_feature_tensor_list = torch.split(
    tensor=all_nodes_feature, 
    split_size_or_sections=number_of_nodes_in_each_graph
)

all_nodes_feature_padded = torch.nn.utils.rnn.pad_sequence(
    sequences=all_nodes_feature_tensor_list, 
    batch_first=True, 
    padding_value=0
)

I hope I could help you with that.
Regards,
Unity05

1 Like

Yup it worked, thanks a lot :slight_smile: