Tokenizer, Dataset, and “collate_fn”
TL;DR
For the best efficiency, override the collate_fn function and leave all tokenization to collate_fn. Don’t run your tokenizer in Dataset!
When you’re training your language model with pytorch, you have to first tokenize the texts before feeding them to the model. Generally speaking, the tokenization can be done at three different places:
You tokenize the texts only once, and then use the results to build the Dataset. This is only feasible when the input texts are small, e.g, a few hundreds.
1
2
3
4
5
| class TextDataset(Dataset):
def __getitem__(self, index):
text = self.text[index]
input_tokens = tokenizer(text, padding="max_length", max_length=128, return_tensors="pt")
return input_tokens
|
Method 2 first does tokenization for every single input (self.text[index]), and then pads every input to the maximum possible length (128 in this case). As you can image, this is very inefficient since you’re not processing the data in batches.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
| class TextDataset(Dataset):
def __getitem__(self, index):
text = self.text[index]
return text
def collate_fn(batch):
# get text
text = batch['text']
# get the input tokens
input_tokens = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
return input_tokens
# init dataloader
dataloader = DataLoader(TextDataloader, batch_size=16, collate_fn=collate_fn)
|
In Method 3, instead of returning the tokenization result (input_tokens) in the __getitem__ function, we simply return the raw text. All the tokenization is done in the collate_fu function. collate_fn is some sort of a data collector: it collects the results of Dataset and then assembles them for the use of DataLoader. So when we run text = batch['text'] in collate_fn, the result is a list of texts whose length is equal to the batch size. Since we only do tokenization for every batch_size inputs, the efficiency has significantly been improved.
Here’s a minimum piece of code to illustrate what I’ve said. The code is based on pytorch-lightning, but you don’t need to be an expert of pytorch-lightning to understand it, since our focus is on the data part, not the model and training part.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
| # The purpose of this file is to test where to put the tokenizer
import os
import torch
from transformers import BertTokenizer, BertModel
from lightning.pytorch import LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
# init the tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
class TextDataset(Dataset):
def __init__(self, text, target):
self.len = len(text)
self.text = text
self.target = target
def __getitem__(self, index):
return (self.text[index], self.target[index])
def __len__(self):
return self.len
def collate_fn(batch):
# unpack batch
text = [_[0] for _ in batch]
target = [_[1] for _ in batch]
# get the input tokens
input_tokens = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
# get the target
target = torch.tensor(target)
return input_tokens, target
class BoringModel(LightningModule):
def __init__(self):
super().__init__()
self.model = BertModel.from_pretrained("bert-base-uncased")
self.linear = torch.nn.Linear(768, 1)
self.loss = torch.nn.MSELoss()
def forward(self, x):
return self.layer(x)
def training_step(self, batch, batch_idx):
# unpack batch
tokens, target = batch
# forward pass
y = self.linear(self.model(**tokens).last_hidden_state).squeeze().sum(-1)
loss = self.loss(y, target)
self.log("train_loss", loss)
return {"loss": loss}
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=1e-3)
def init_inputs(n):
# train/val/test: 7:1:2
text = ["Yes!", "A sentence with more words.", "Come on!"] * n
target = [1.0, 2, 3] * n
return text, target
def run():
# hypyterparameters
batch_size = 32
# initialize the text
train_text, train_target = init_inputs(n=int(1e4))
# init the dataset
train_data = DataLoader(
TextDataset(train_text, train_target),
batch_size=batch_size,
collate_fn=collate_fn,
drop_last=True,
)
model = BoringModel()
trainer = Trainer(
accelerator="gpu",
devices=2,
strategy="ddp_find_unused_parameters_true",
default_root_dir=os.getcwd(),
num_sanity_val_steps=0,
max_epochs=1,
enable_model_summary=True,
)
trainer.fit(model, train_dataloaders=train_data)
if __name__ == "__main__":
run()
|