I'm trying to implement a model using lstm to generate sonnects in pytorch. When I tested the dataLoader, It took a lot of time and could not return data as expected. I review it and don't know where error locates. Please help me, here is the code(more detail in(https://paste.ubuntu.com/p/vspS3msNVW/)) (10th cell)
class SonnetDataset(Dataset):
def __init__(self, sonnet_in_ids: list, vocab: list, max_seq_length: int):
super().__init__()
self.data = sonnet_in_ids
self.vocab = vocab
self.vocab_size = len(vocab)
self.pad_id = self.vocab.index('<PAD>')
self.start_id = self.vocab.index('<START>')
self.end_id = self.vocab.index('<END>')
self.max_seq_length = max_seq_length + 2
print('init successfully')
def __len__(self):
return len(self.data)
def __getitem__(self, index) -> torch.LongTensor:
print('get item')
x = self.data[index]
x = [self.start_id] + x + [self.end_id]
# padding
x += [self.pad_id] * (self.max_seq_length - len(x))
x = torch.LongTensor(x)
print(x)
return x
batch_size = 4
train_set = SonnetDataset(sonnet_in_ids=sonnets_in_ids, vocab=vocab, max_seq_length=max_length)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True)
next(iter(train_loader))
Could not get data