【Pytorch】collate_fn函数

【参考：pytorch中collate_fn函数的使用&如何向collate_fn函数传参_XJTU-Qidong的博客-CSDN博客】

collate_fn

class MyDataset(Dataset):
    def __init__(self, datas, tags, word2idx, tag2idx):
        self.datas = datas
        self.tags = tags
        self.word2idx = word2idx
        self.tag2idx = tag2idx

    def __getitem__(self, index):
        data = self.datas[index]
        tag = self.tags[index]
        # 如果word2idx中没有该词就用UNK代替
        data_index = [self.word2idx.get(i, self.word2idx["<UNK>"]) for i in data]  # 获取每句话所有字的索引
        tag_index = [self.tag2idx[i] for i in tag]

        return data_index, tag_index

    def __len__(self):
        return len(self.datas)

    def pro_batch_data(self, batch_datas): 
    	# batch_datas 是一个批次的数据 List类型
        # 因为每句话的长度不一样，拼接数据维度不一致会导致错误，所以自己手动拼接 未完全理解
        datas = []
        tags = []
        batch_lens = []

        for data, tag in batch_datas:
            datas.append(data)
            tags.append(tag)
            batch_lens.append(len(data))
        batch_max_len = max(batch_lens)

        # 填充至相同长度
        datas = [i + [self.word2idx['<PAD>']] * (batch_max_len - len(i)) for i in datas]  # i是，每句话
        tags = [i + [self.tag2idx['<PAD>']] * (batch_max_len - len(i)) for i in tags]

        # return torch.IntTensor(datas), torch.LongTensor(tags)
        return torch.tensor(datas, dtype=torch.int64), torch.tensor(tags, dtype=torch.long)


train_dataset = MyDataset(train_data, train_tag, word2id, tag2id)
train_dataloader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=False
                                  , collate_fn=train_dataset.pro_batch_data)

也可以写在外部

def collate_fn(batch_data,word2idx):
    # batch_data 是一个批次的数据 List
    # 因为每句话的长度不一样，拼接数据维度不一致会导致错误，所以自己手动截断补全
    datas = []
    tags = []
    batch_lens = []

    for data, tag in batch_data:
        datas.append(data)
        tags.append(tag)
        batch_lens.append(len(data))
    batch_max_len = max(batch_lens)

    # 填充至相同长度
    datas = [i + [word2idx['<PAD>']] * (batch_max_len - len(i)) for i in datas]  # i是，每句话

    return torch.tensor(datas),torch.tensor(tags) # 转成tensor


batch_size = 8

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
                          collate_fn=lambda x: collate_fn(x, word2idx)) # x是一个批次的数据
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True,
                         collate_fn=lambda x: collate_fn(x, word2idx))

device = 'cuda' if torch.cuda.is_available() else 'cpu'