使用 pytorch的dataloader

from torch.utils.data import DataLoader
from datasets import 



datasets = {
        'train': Hijackdataset(
            input_ids=token_ids,
            lens=lens,
            mask=mask,
            labels=labels,
            task=task
        ),
        'test': Hijackdataset(
            input_ids=test_token_ids,
            lens=test_lens,
            mask=test_mask,
            labels=test_labels,
            task=task
        )
    }

dataloaders = {
        'train': DataLoader(
            dataset=datasets['train'],
            batch_size=bs,
        ),
        'test': DataLoader(
            dataset=datasets['test'], 
            batch_size=bs
        )

定义datasets

class Hijackdataset(Dataset):
    def __init__(self, input_ids, lens, mask, labels, task):
        self.input_ids = torch.tensor(input_ids)
        self.lens = lens
        self.mask = torch.tensor(mask, dtype=torch.float32)
        self.labels = labels

    def __len__(self):
        return self.labels['a'].shape[0]

    def __getitem__(self, idx):
        input = self.input_ids[idx]
        mask = self.mask[idx]
        length = self.lens[idx]
        #
        label_A = self.label['a']
        label_B = self.label['b']
        label_C = self.label['c']
        return input, length, mask, label_A, label_B, label_C

实现dataloader的类，一定要有len和getitem

for inputs, lens, mask, label_A, label_B, label_C in tqdm(dataloader, desc='Training M'):
            iters_per_epoch += 1
            inputs = inputs.to(device=self.device)
            lens = lens.to(device=self.device)
            mask = mask.to(device=self.device)
            label_A = label_A.to(device=self.device)
            label_B = label_B.to(device=self.device)
            label_C = label_C.to(device=self.device)
            self.optimizer.zero_grad()
            with torch.set_grad_enabled(True):
                all_logits = self.model(inputs, lens, mask)
                y_pred_A = all_logits[0].argmax(dim=1).cpu().numpy()
                y_pred_B = all_logits[1][:, 0:2].argmax(dim=1)
                y_pred_C = all_logits[2][:, 0:3].argmax(dim=1)
                _loss = self.loss_weights[0] * self.criterion(all_logits[0], label_A)
                _loss += self.loss_weights[1] * self.criterion(all_logits[1], label_B)
                _loss += self.loss_weights[2] * self.criterion(all_logits[2], label_C)
                loss += _loss.item()
                # Backward
                _loss.backward()

这里inputs lens mask label对应 getitem中的属性（在第一次使用时初始化，在实例中for inputs in dataloader时初始化）