Pytorch 使用不同的张量长度

Question

I am working with word embeddings, and each phrase has different length. My dataset contains a list of vectors, one vector of size 300 for each word in the phrase. Maybe one phrase has 20 words, maybe it has 2, etc.

For instance:

X.iloc[0:10, ]
0    [[-0.51389, -0.55286, -0.28648, -0.18608, -0.0...
1    [[0.33837, -0.13933, -0.096114, 0.40103, 0.041...
2    [[-0.078564, 0.18702, -0.35168, 0.067557, 0.11...
3    [[0.047356, -0.10216, -0.15738, -0.04521, 0.26...
4    [[0.16781, -0.31793, -0.21877, 0.28025, 0.3364...
5    [[-0.4509, 0.077681, -0.058347, 0.2859, -0.369...
6    [[0.018223, -0.012323, 0.035569, 0.24232, -0.1...
7    [[-0.19265, 0.45863, -0.33841, -0.16293, -0.26...
8    [[0.10751, 0.15958, 0.13332, 0.16642, -0.03273...
9    [[0.35259, 0.60833, 0.051335, -0.079285, -0.35...
Name: embedding, dtype: object

len(X.iloc[0])
313
len(X.iloc[1])
2

The targets are just a numeric integer, from 0 to 5.

How can I pad this sequence using pytorch to feed a neural network of fixed size? I saw something with collate_fn , however I think it only works for batches, and not for the whole dataset.

Answer 1

It seems correct that collate_fn can help in this case. Indeed, it is applied on batch level, but iterating through all batches will process entire dataset. For padding you can use some of the following functions:

torch.nn.functional.pad

torch.nn.utils.rnn.pad_sequence

Doing it manually in a custom way is also an option.

Here is toy example with torch.nn.functional.pad and random dataset:

import torch
import random
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset


MAX_PHRASE_LENGTH = 20


class PhraseDataset(Dataset):

    def __len__(self):
        return 5

    def __getitem__(self, item):
        phrase = torch.rand((random.randint(2, MAX_PHRASE_LENGTH), 300), dtype=torch.float)
        target = torch.tensor(random.randint(0, 5), dtype=torch.float)

        return phrase, target


def collate_fn(batch):
    targets = []
    phrases = []
    for (phrase, target) in batch:
        targets.append(target)
        padding = MAX_PHRASE_LENGTH - phrase.size()[0]
        if padding > 0:
            phrases.append(F.pad(phrase, (0, 0, padding, 0), "constant", 0))
        else:
            phrases.append(phrase)

    return torch.stack(phrases), torch.tensor(targets, dtype=torch.float)


dataset = PhraseDataset()
data_loader = DataLoader(dataset, batch_size=5, collate_fn=collate_fn)

for phrases, targets in data_loader:
    print('Phrases size:', phrases.size())

Output:

Phrases size: torch.Size([5, 20, 300])

Although padding could actually be applied in a dataset, before collate_fn is called, so collate_fn won't be needed in this case at all.