【参考: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'