pytorch Dataset类

官网：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