官网：Dataset类
简介
Dataset是一个包装类，用来将数据包装为Dataset类，然后传入DataLoader中，我们再使用DataLoader这个类来更加快捷的对数据进行操作。在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。直至把所有的数据都抛出。就是做一个数据的初始化。
官网参数
参数解释：
dataset：需要load的数据
batch_size：每个batch的大小，一次返回几条数据处理
shuffle：是否进行shuffle操作
num_workers：加载数据的时候使用几个子进程
使用方式很固定，直接记住代码格式：
import torch.utils.data as Data
input_batch, target_batch = make_data(sentence)
dataset = Data.TensorDataset(input_batch, target_batch)
loader = Data.DataLoader(dataset, batch_size=16, shuffle=True)
#训练
for epoch in range(10000):
for x, y in loader:
pred = model(x)
loss = criterion(pred, y)
if (epoch + 1) % 1000 == 0:
print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))
optimizer.zero_grad()#torch三件套
loss.backward()
optimizer.step()
####分割线就是我
#将数据处理成需要的格式tensor类型
train_inputs, train_token, train_mask, train_labels=make_data(sentence)
train_data = Data.TensorDataset(train_inputs, train_token, train_mask, train_labels)
train_dataloader = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
for _ in range(2):
for i, batch in enumerate(train_dataloader):#因为数据是四种，所以用迭代器的方式处理，bach=[tensor1,tensor2,tensor3，tensor4]
batch = tuple(t.to(device) for t in batch)#转换成tuple并且指定cuda运行
loss = model(batch[0], token_type_ids=batch[1], attention_mask=batch[2], labels=batch[3])[0]
print(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 10 == 0:
eval(model, validation_dataloader)
以上图第二段例子代码来看：
首先是把原句子输入预处理函数 make_data；这个函数是自己定义的，按照任务要求，一般调用 tokenizer 来处理；举个例子：
下面是一个预处理的函数示例：
def make_data(sentences,labels):
input_ids,token_type_ids,attention_mask=[],[],[]
#input_ids是每个词对应的索引idx ;token_type_ids是对应的0和1，标识是第几个句子；attention_mask是对句子长度做pad
#input_ids=[22,21,...499] token_type_ids=[0,0,0,0,1,1,1,1] ;attention_mask=[1,1,1,1,1,0,0,0]补零
for i in range(len(sentences)):
encoded_dict = tokenizer.encode_plus(
sentences[i],
add_special_tokens = True, # 添加 '[CLS]' 和 '[SEP]'
max_length = 96, # 填充 & 截断长度
pad_to_max_length = True,
return_tensors = 'pt', # 返回 pytorch tensors 格式的数据
)
input_ids.append(encoded_dict['input_ids'])
token_type_ids.append(encoded_dict['token_type_ids'])
attention_mask.append(encoded_dict['attention_mask'])
input_ids = torch.LongTensor(input_ids)#每个词对应的索引
token_type_ids = torch.LongTensor(token_type_ids)#0&1标识是哪个句子
attention_mask = torch.LongTensor(attention_mask)#[11100]padding之后的句子
labels = torch.LongTensor(labels)#所有实例的label对应的索引idx
return input_ids, token_type_ids, attention_mask, labels
经过这个处理之后的返回值 然后调用 Data.TensorDataset 和 Data.DataLoader；这个其实就是把所有需要的数据都装载到了一个dict 里面，然后被封装好，设定 batch 的大小，然后进行训练：
train_inputs, train_token, train_mask, train_labels=make_data(sentence)
train_data = Data.TensorDataset(train_inputs, train_token, train_mask, train_labels)
train_dataloader = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
for i, batch in enumerate(train_dataloader):
batch = tuple(t.to(device) for t in batch)#转换成tuple并且指定cuda运行
loss = model(batch[0], token_type_ids=batch[1], attention_mask=batch[2], labels=batch[3])[0]
#这里通过batch下标就可以达到访问各个类别训练数据的目的
改写dataset 类
在我们自己的模型训练中，常常需要使用非官方自制的数据集。
我们可以通过改写torch.utils.data.Dataset中的getitem和len来载入我们自己的数据集。
getitem：获取数据集中的数据
len：获取整个数据集的长度（即个数）
看下面的一个例子：
class IMDbDataset(torch.utils.data.Dataset):
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item['labels'] = torch.tensor(self.labels[idx])
return item
#item={"input_ids":[] ,"token_type_ids":[],"attention_mask":[],"labels":[]}
def __len__(self):
return len(self.labels)
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)
train_dataset = IMDbDataset(train_encodings, train_labels)
test_dataset = IMDbDataset(test_encodings, test_labels)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
for epoch in range(3):
for batch in train_loader:
optim.zero_grad()
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs[0]
loss.backward()
optim.step()
这个例子跟上文例子其实相差不多，只是这个地方把 作为分类任务的 label 数据传入 重写的 dataset 类，加入训练数据的 dict 中
item={"input_ids":[] ,"token_type_ids":[],"attention_mask":[],"labels":[]}
上文的例子是直接把 label 和其他类别的训练数据如：input_ids ，attention_mask 等一起传入 Data.TensorDataset() 其实也是一样的目的，这个地方只是给一个简单的例子，表示在实际的任务中，如果需要更加复杂的封装训练数据，需要重写 dataset 类中的这些函数来达到加载数据的目的；
参数 collate_fn 用法：
将一个list的sample组成一个mini-batch的函数，可以自定义这个参数，将自己的数据处理成一个 batch，padding mask 等操作：
def custom_collate(batch):
transposed = list(zip(*batch))
lst = []
# transposed[0]: list of token ids of text
padded_seq = []
max_seq_len = len(max(transposed[0], key=len))
for seq in transposed[0]:
padded_seq.append(seq + [0] * (max_seq_len - len(seq)))
lst.append(torch.LongTensor(padded_seq))
# tansposed[1]: list of tag ids of SAME LENGTH!
padded_tag = []
att_mask = []
for seq in transposed[1]:
padded_tag.append(seq + [0] * (max_seq_len - len(seq)))
att_mask.append([1] * len(seq) + [0] * (max_seq_len - len(seq)))
lst.append(torch.LongTensor(padded_tag))
lst.append(torch.FloatTensor(att_mask))
return lst
dataset = MyDataset(train_tokens, train_tags)
train_loader = DataLoader(dataset=dataset, batch_size=train_batch, collate_fn=custom_collate, shuffle=True)
- simple example:
# a simple custom collate function, just to show the idea
def my_collate(batch):
data = [item[0] for item in batch]
target = [item[1] for item in batch]
target = torch.LongTensor(target)
return [data, target]
1.
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
"""定义dataloader,在训练阶段shuffle数据，预测阶段不需要shuffle"""
train_data = TensorDataset(tr_inputs, tr_masks, tr_tags)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)
valid_data = TensorDataset(val_inputs, val_masks, val_tags)
valid_sampler = SequentialSampler(valid_data)
valid_dataloader = DataLoader(valid_data, sampler=valid_sampler, batch_size=batch_size)
参考：
huggingface fine-tuning with custom dataset
BERT Fine-Tuning Tutorial with PyTorch · Chris McCormick (mccormickml.com)
https://ptorch.com/docs/1/utils-data
