Pytorch with CUDA throws RuntimeError when using pack_padded_sequence
Question:
I am trying to train a BiLSTM-CRF on detecting new NER entities with Pytorch.
To do so, I am using a snippet of code derivated from the Pytorch Advanced tutorial. This snippet implements batch training.
I followed the READ-ME in order to present data as required. Everything works great on CPU, but when I’m trying to get it to GPU, the following error occur :
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-23-794982510db6> in <module>
4 batch_input, batch_input_lens, batch_mask, batch_target = batch_info
5
----> 6 loss_train = model.neg_log_likelihood(batch_input, batch_input_lens, batch_mask, batch_target)
7 optimizer.zero_grad()
8 loss_train.backward()
<ipython-input-11-e44ffbf7d75f> in neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target)
185
186 def neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target):
--> 187 feats = self.bilstm(batch_input, batch_input_lens, batch_mask)
188 gold_score = self.CRF.score_sentence(feats, batch_target)
189 forward_score = self.CRF.score_z(feats, batch_input_lens)
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1050 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051 return forward_call(*input, **kwargs)
1052 # Do not call functions when jit is used
1053 full_backward_hooks, non_full_backward_hooks = [], []
<ipython-input-11-e44ffbf7d75f> in forward(self, batch_input, batch_input_lens, batch_mask)
46 batch_input = self.word_embeds(batch_input) # size: #batch * padding_length * embedding_dim
47 batch_input = rnn_utils.pack_padded_sequence(
---> 48 batch_input, batch_input_lens, batch_first=True)
49 batch_output, self.hidden = self.lstm(batch_input, self.hidden)
50 self.repackage_hidden(self.hidden)
/opt/conda/lib/python3.7/site-packages/torch/nn/utils/rnn.py in pack_padded_sequence(input, lengths, batch_first, enforce_sorted)
247
248 data, batch_sizes =
--> 249 _VF._pack_padded_sequence(input, lengths, batch_first)
250 return _packed_sequence_init(data, batch_sizes, sorted_indices, None)
251
RuntimeError: 'lengths' argument should be a 1D CPU int64 tensor, but got 1D cuda:0 Long tensor`
If I understand well, pack_padded_sequence need the tensor to be on CPU rather than GPU. Unfortunately the pack_padded_sequence is called by my forward function and I can’t see any way to do so without going back to CPU for the whole training.
Here is the complete code.
Classes definitions :
import torch
import torch.nn as nn
import torch.nn.utils.rnn as rnn_utils
class BiLSTM(nn.Module):
def __init__(self, vocab_size, tagset, embedding_dim, hidden_dim,
num_layers, bidirectional, dropout, pretrained=None):
super(BiLSTM, self).__init__()
self.embedding_dim = embedding_dim
self.hidden_dim = hidden_dim
self.tagset_size = len(tagset)
self.bidirectional = bidirectional
self.num_layers = num_layers
self.word_embeds = nn.Embedding(vocab_size+2, embedding_dim)
if pretrained is not None:
self.word_embeds = nn.Embedding.from_pretrained(pretrained)
self.lstm = nn.LSTM(
input_size=embedding_dim,
hidden_size=hidden_dim // 2 if bidirectional else hidden_dim,
num_layers=num_layers,
dropout=dropout,
bidirectional=bidirectional,
batch_first=True,
)
self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
self.hidden = None
def init_hidden(self, batch_size, device):
init_hidden_dim = self.hidden_dim // 2 if self.bidirectional else self.hidden_dim
init_first_dim = self.num_layers * 2 if self.bidirectional else self.num_layers
self.hidden = (
torch.randn(init_first_dim, batch_size, init_hidden_dim).to(device),
torch.randn(init_first_dim, batch_size, init_hidden_dim).to(device)
)
def repackage_hidden(self, hidden):
"""Wraps hidden states in new Tensors, to detach them from their history."""
if isinstance(hidden, torch.Tensor):
return hidden.detach_().to(device)
else:
return tuple(self.repackage_hidden(h) for h in hidden)
def forward(self, batch_input, batch_input_lens, batch_mask):
batch_size, padding_length = batch_input.size()
batch_input = self.word_embeds(batch_input) # size: #batch * padding_length * embedding_dim
batch_input = rnn_utils.pack_padded_sequence(
batch_input, batch_input_lens, batch_first=True)
batch_output, self.hidden = self.lstm(batch_input, self.hidden)
self.repackage_hidden(self.hidden)
batch_output, _ = rnn_utils.pad_packed_sequence(batch_output, batch_first=True)
batch_output = batch_output.contiguous().view(batch_size * padding_length, -1)
batch_output = batch_output[batch_mask, ...]
out = self.hidden2tag(batch_output)
return out
def neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target):
loss = nn.CrossEntropyLoss(reduction='mean')
feats = self(batch_input, batch_input_lens, batch_mask)
batch_target = torch.cat(batch_target, 0).to(device)
return loss(feats, batch_target)
def predict(self, batch_input, batch_input_lens, batch_mask):
feats = self(batch_input, batch_input_lens, batch_mask)
val, pred = torch.max(feats, 1)
return pred
class CRF(nn.Module):
def __init__(self, tagset, start_tag, end_tag, device):
super(CRF, self).__init__()
self.tagset_size = len(tagset)
self.START_TAG_IDX = tagset.index(start_tag)
self.END_TAG_IDX = tagset.index(end_tag)
self.START_TAG_TENSOR = torch.LongTensor([self.START_TAG_IDX]).to(device)
self.END_TAG_TENSOR = torch.LongTensor([self.END_TAG_IDX]).to(device)
# trans: (tagset_size, tagset_size) trans (i, j) means state_i -> state_j
self.trans = nn.Parameter(
torch.randn(self.tagset_size, self.tagset_size)
)
# self.trans.data[...] = 1
self.trans.data[:, self.START_TAG_IDX] = -10000
self.trans.data[self.END_TAG_IDX, :] = -10000
self.device = device
def init_alpha(self, batch_size, tagset_size):
return torch.full((batch_size, tagset_size, 1), -10000, dtype=torch.float, device=self.device)
def init_path(self, size_shape):
# Initialization Path - LongTensor + Device + Full_value=0
return torch.full(size_shape, 0, dtype=torch.long, device=self.device)
def _iter_legal_batch(self, batch_input_lens, reverse=False):
index = torch.arange(0, batch_input_lens.sum(), dtype=torch.long)
packed_index = rnn_utils.pack_sequence(
torch.split(index, batch_input_lens.tolist())
)
batch_iter = torch.split(packed_index.data, packed_index.batch_sizes.tolist())
batch_iter = reversed(batch_iter) if reverse else batch_iter
for idx in batch_iter:
yield idx, idx.size()[0]
def score_z(self, feats, batch_input_lens):
# 模拟packed pad过程
tagset_size = feats.shape[1]
batch_size = len(batch_input_lens)
alpha = self.init_alpha(batch_size, tagset_size)
alpha[:, self.START_TAG_IDX, :] = 0 # Initialization
for legal_idx, legal_batch_size in self._iter_legal_batch(batch_input_lens):
feat = feats[legal_idx, ].view(legal_batch_size, 1, tagset_size) #
# #batch * 1 * |tag| + #batch * |tag| * 1 + |tag| * |tag| = #batch * |tag| * |tag|
legal_batch_score = feat + alpha[:legal_batch_size, ] + self.trans
alpha_new = torch.logsumexp(legal_batch_score, 1).unsqueeze(2).to(device)
alpha[:legal_batch_size, ] = alpha_new
alpha = alpha + self.trans[:, self.END_TAG_IDX].unsqueeze(1)
score = torch.logsumexp(alpha, 1).sum().to(device)
return score
def score_sentence(self, feats, batch_target):
# CRF Batched Sentence Score
# feats: (#batch_state(#words), tagset_size)
# batch_target: list<torch.LongTensor> At least One LongTensor
# Warning: words order = batch_target order
def _add_start_tag(target):
return torch.cat([self.START_TAG_TENSOR, target]).to(device)
def _add_end_tag(target):
return torch.cat([target, self.END_TAG_TENSOR]).to(device)
from_state = [_add_start_tag(target) for target in batch_target]
to_state = [_add_end_tag(target) for target in batch_target]
from_state = torch.cat(from_state).to(device)
to_state = torch.cat(to_state).to(device)
trans_score = self.trans[from_state, to_state]
gather_target = torch.cat(batch_target).view(-1, 1).to(device)
emit_score = torch.gather(feats, 1, gather_target).to(device)
return trans_score.sum() + emit_score.sum()
def viterbi(self, feats, batch_input_lens):
word_size, tagset_size = feats.shape
batch_size = len(batch_input_lens)
viterbi_path = self.init_path(feats.shape) # use feats.shape to init path.shape
alpha = self.init_alpha(batch_size, tagset_size)
alpha[:, self.START_TAG_IDX, :] = 0 # Initialization
for legal_idx, legal_batch_size in self._iter_legal_batch(batch_input_lens):
feat = feats[legal_idx, :].view(legal_batch_size, 1, tagset_size)
legal_batch_score = feat + alpha[:legal_batch_size, ] + self.trans
alpha_new, best_tag = torch.max(legal_batch_score, 1).to(device)
alpha[:legal_batch_size, ] = alpha_new.unsqueeze(2)
viterbi_path[legal_idx, ] = best_tag
alpha = alpha + self.trans[:, self.END_TAG_IDX].unsqueeze(1)
path_score, best_tag = torch.max(alpha, 1).to(device)
path_score = path_score.squeeze() # path_score=#batch
best_paths = self.init_path((word_size, 1))
for legal_idx, legal_batch_size in self._iter_legal_batch(batch_input_lens, reverse=True):
best_paths[legal_idx, ] = best_tag[:legal_batch_size, ] #
backword_path = viterbi_path[legal_idx, ] # 1 * |Tag|
this_tag = best_tag[:legal_batch_size, ] # 1 * |legal_batch_size|
backword_tag = torch.gather(backword_path, 1, this_tag).to(device)
best_tag[:legal_batch_size, ] = backword_tag
# never computing <START>
# best_paths = #words
return path_score.view(-1), best_paths.view(-1)
class BiLSTM_CRF(nn.Module):
def __init__(self, vocab_size, tagset, embedding_dim, hidden_dim,
num_layers, bidirectional, dropout, start_tag, end_tag, device, pretrained=None):
super(BiLSTM_CRF, self).__init__()
self.bilstm = BiLSTM(vocab_size, tagset, embedding_dim, hidden_dim,
num_layers, bidirectional, dropout, pretrained)
self.CRF = CRF(tagset, start_tag, end_tag, device)
def init_hidden(self, batch_size, device):
self.bilstm.hidden = self.bilstm.init_hidden(batch_size, device)
def forward(self, batch_input, batch_input_lens, batch_mask):
feats = self.bilstm(batch_input, batch_input_lens, batch_mask)
score, path = self.CRF.viterbi(feats, batch_input_lens)
return path
def neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target):
feats = self.bilstm(batch_input, batch_input_lens, batch_mask)
gold_score = self.CRF.score_sentence(feats, batch_target)
forward_score = self.CRF.score_z(feats, batch_input_lens)
return forward_score - gold_score
def predict(self, batch_input, batch_input_lens, batch_mask):
return self(batch_input, batch_input_lens, batch_mask)
Training cell :
def prepare_sequence(seq, to_ix, device):
idxs = [to_ix[w] for w in seq]
return torch.tensor(idxs, dtype=torch.long).to(device)
def prepare_labels(lab, tag_to_ix, device):
idxs = [tag_to_ix[w] for w in lab]
return torch.tensor(idxs, dtype=torch.long).to(device)
class PadSequence:
def __call__(self, batch):
device = torch.device('cuda')
# Let's assume that each element in "batch" is a tuple (data, label).
# Sort the batch in the descending order
sorted_batch = sorted(batch, key=lambda x: len(x[0]), reverse=True)
# Get each sequence and pad it
sequences = [x[0] for x in sorted_batch]
sentence_in =[prepare_sequence(x, word_to_ix, device) for x in sequences]
sequences_padded = torch.nn.utils.rnn.pad_sequence(sentence_in, padding_value = len(word_to_ix) +1, batch_first=True).to(device)
lengths = torch.LongTensor([len(x) for x in sequences]).to(device)
masks = [True if index_word!=len(word_to_ix)+1 else False for sentence in sequences_padded for index_word in sentence ]
labels = [x[1] for x in sorted_batch]
labels_in = [prepare_sequence(x, tag_to_ix, device) for x in labels]
return sequences_padded, lengths, masks, labels_in
{ .... code to get the data formatted...}
device = torch.device("cuda")
batch_size = 64
START_TAG = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 200
HIDDEN_DIM = 20
NUM_LAYER = 3
BIDIRECTIONNAL = True
DROPOUT = 0.1
train_iter = DataLoader(dataset=training_data, collate_fn=PadSequence(), batch_size=64, shuffle=True)
model = BiLSTM_CRF(len(word_to_ix), tagset, EMBEDDING_DIM, HIDDEN_DIM, NUM_LAYER, BIDIRECTIONNAL, DROPOUT, START_TAG, STOP_TAG, device ).to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)
model.init_hidden(batch_size, device)
with tqdm(total=len(train_iter)) as progress_bar:
for batch_info in train_iter:
batch_input, batch_input_lens, batch_mask, batch_target = batch_info
loss_train = model.neg_log_likelihood(batch_input, batch_input_lens, batch_mask, batch_target)
optimizer.zero_grad()
loss_train.backward()
optimizer.step()
progress_bar.update(1) # update progress
Answers:
Within PadSequence function (which acts as a collate_fn which gathers samples and makes a batch from them) you are explicitly casting to cuda device, namely:
class PadSequence:
def __call__(self, batch):
device = torch.device('cuda')
# Left rest of the code for brevity
...
lengths = torch.LongTensor([len(x) for x in sequences]).to(device)
...
return sequences_padded, lengths, masks, labels_in
You don’t need to cast your data when creating batch, we usually do that right before pushing the examples through neural network.
Also you should at least define the device like this:
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
or even better leave the choice of device for you/user in some part of the code where you setup everything.
The error lays in
lengths = torch.LongTensor([len(x) for x in sequences]).to(device)
change to
lengths = torch.LongTensor([len(x) for x in sequences]).cpu()
I am trying to train a BiLSTM-CRF on detecting new NER entities with Pytorch.
To do so, I am using a snippet of code derivated from the Pytorch Advanced tutorial. This snippet implements batch training.
I followed the READ-ME in order to present data as required. Everything works great on CPU, but when I’m trying to get it to GPU, the following error occur :
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-23-794982510db6> in <module>
4 batch_input, batch_input_lens, batch_mask, batch_target = batch_info
5
----> 6 loss_train = model.neg_log_likelihood(batch_input, batch_input_lens, batch_mask, batch_target)
7 optimizer.zero_grad()
8 loss_train.backward()
<ipython-input-11-e44ffbf7d75f> in neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target)
185
186 def neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target):
--> 187 feats = self.bilstm(batch_input, batch_input_lens, batch_mask)
188 gold_score = self.CRF.score_sentence(feats, batch_target)
189 forward_score = self.CRF.score_z(feats, batch_input_lens)
/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1050 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051 return forward_call(*input, **kwargs)
1052 # Do not call functions when jit is used
1053 full_backward_hooks, non_full_backward_hooks = [], []
<ipython-input-11-e44ffbf7d75f> in forward(self, batch_input, batch_input_lens, batch_mask)
46 batch_input = self.word_embeds(batch_input) # size: #batch * padding_length * embedding_dim
47 batch_input = rnn_utils.pack_padded_sequence(
---> 48 batch_input, batch_input_lens, batch_first=True)
49 batch_output, self.hidden = self.lstm(batch_input, self.hidden)
50 self.repackage_hidden(self.hidden)
/opt/conda/lib/python3.7/site-packages/torch/nn/utils/rnn.py in pack_padded_sequence(input, lengths, batch_first, enforce_sorted)
247
248 data, batch_sizes =
--> 249 _VF._pack_padded_sequence(input, lengths, batch_first)
250 return _packed_sequence_init(data, batch_sizes, sorted_indices, None)
251
RuntimeError: 'lengths' argument should be a 1D CPU int64 tensor, but got 1D cuda:0 Long tensor`
If I understand well, pack_padded_sequence need the tensor to be on CPU rather than GPU. Unfortunately the pack_padded_sequence is called by my forward function and I can’t see any way to do so without going back to CPU for the whole training.
Here is the complete code.
Classes definitions :
import torch
import torch.nn as nn
import torch.nn.utils.rnn as rnn_utils
class BiLSTM(nn.Module):
def __init__(self, vocab_size, tagset, embedding_dim, hidden_dim,
num_layers, bidirectional, dropout, pretrained=None):
super(BiLSTM, self).__init__()
self.embedding_dim = embedding_dim
self.hidden_dim = hidden_dim
self.tagset_size = len(tagset)
self.bidirectional = bidirectional
self.num_layers = num_layers
self.word_embeds = nn.Embedding(vocab_size+2, embedding_dim)
if pretrained is not None:
self.word_embeds = nn.Embedding.from_pretrained(pretrained)
self.lstm = nn.LSTM(
input_size=embedding_dim,
hidden_size=hidden_dim // 2 if bidirectional else hidden_dim,
num_layers=num_layers,
dropout=dropout,
bidirectional=bidirectional,
batch_first=True,
)
self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
self.hidden = None
def init_hidden(self, batch_size, device):
init_hidden_dim = self.hidden_dim // 2 if self.bidirectional else self.hidden_dim
init_first_dim = self.num_layers * 2 if self.bidirectional else self.num_layers
self.hidden = (
torch.randn(init_first_dim, batch_size, init_hidden_dim).to(device),
torch.randn(init_first_dim, batch_size, init_hidden_dim).to(device)
)
def repackage_hidden(self, hidden):
"""Wraps hidden states in new Tensors, to detach them from their history."""
if isinstance(hidden, torch.Tensor):
return hidden.detach_().to(device)
else:
return tuple(self.repackage_hidden(h) for h in hidden)
def forward(self, batch_input, batch_input_lens, batch_mask):
batch_size, padding_length = batch_input.size()
batch_input = self.word_embeds(batch_input) # size: #batch * padding_length * embedding_dim
batch_input = rnn_utils.pack_padded_sequence(
batch_input, batch_input_lens, batch_first=True)
batch_output, self.hidden = self.lstm(batch_input, self.hidden)
self.repackage_hidden(self.hidden)
batch_output, _ = rnn_utils.pad_packed_sequence(batch_output, batch_first=True)
batch_output = batch_output.contiguous().view(batch_size * padding_length, -1)
batch_output = batch_output[batch_mask, ...]
out = self.hidden2tag(batch_output)
return out
def neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target):
loss = nn.CrossEntropyLoss(reduction='mean')
feats = self(batch_input, batch_input_lens, batch_mask)
batch_target = torch.cat(batch_target, 0).to(device)
return loss(feats, batch_target)
def predict(self, batch_input, batch_input_lens, batch_mask):
feats = self(batch_input, batch_input_lens, batch_mask)
val, pred = torch.max(feats, 1)
return pred
class CRF(nn.Module):
def __init__(self, tagset, start_tag, end_tag, device):
super(CRF, self).__init__()
self.tagset_size = len(tagset)
self.START_TAG_IDX = tagset.index(start_tag)
self.END_TAG_IDX = tagset.index(end_tag)
self.START_TAG_TENSOR = torch.LongTensor([self.START_TAG_IDX]).to(device)
self.END_TAG_TENSOR = torch.LongTensor([self.END_TAG_IDX]).to(device)
# trans: (tagset_size, tagset_size) trans (i, j) means state_i -> state_j
self.trans = nn.Parameter(
torch.randn(self.tagset_size, self.tagset_size)
)
# self.trans.data[...] = 1
self.trans.data[:, self.START_TAG_IDX] = -10000
self.trans.data[self.END_TAG_IDX, :] = -10000
self.device = device
def init_alpha(self, batch_size, tagset_size):
return torch.full((batch_size, tagset_size, 1), -10000, dtype=torch.float, device=self.device)
def init_path(self, size_shape):
# Initialization Path - LongTensor + Device + Full_value=0
return torch.full(size_shape, 0, dtype=torch.long, device=self.device)
def _iter_legal_batch(self, batch_input_lens, reverse=False):
index = torch.arange(0, batch_input_lens.sum(), dtype=torch.long)
packed_index = rnn_utils.pack_sequence(
torch.split(index, batch_input_lens.tolist())
)
batch_iter = torch.split(packed_index.data, packed_index.batch_sizes.tolist())
batch_iter = reversed(batch_iter) if reverse else batch_iter
for idx in batch_iter:
yield idx, idx.size()[0]
def score_z(self, feats, batch_input_lens):
# 模拟packed pad过程
tagset_size = feats.shape[1]
batch_size = len(batch_input_lens)
alpha = self.init_alpha(batch_size, tagset_size)
alpha[:, self.START_TAG_IDX, :] = 0 # Initialization
for legal_idx, legal_batch_size in self._iter_legal_batch(batch_input_lens):
feat = feats[legal_idx, ].view(legal_batch_size, 1, tagset_size) #
# #batch * 1 * |tag| + #batch * |tag| * 1 + |tag| * |tag| = #batch * |tag| * |tag|
legal_batch_score = feat + alpha[:legal_batch_size, ] + self.trans
alpha_new = torch.logsumexp(legal_batch_score, 1).unsqueeze(2).to(device)
alpha[:legal_batch_size, ] = alpha_new
alpha = alpha + self.trans[:, self.END_TAG_IDX].unsqueeze(1)
score = torch.logsumexp(alpha, 1).sum().to(device)
return score
def score_sentence(self, feats, batch_target):
# CRF Batched Sentence Score
# feats: (#batch_state(#words), tagset_size)
# batch_target: list<torch.LongTensor> At least One LongTensor
# Warning: words order = batch_target order
def _add_start_tag(target):
return torch.cat([self.START_TAG_TENSOR, target]).to(device)
def _add_end_tag(target):
return torch.cat([target, self.END_TAG_TENSOR]).to(device)
from_state = [_add_start_tag(target) for target in batch_target]
to_state = [_add_end_tag(target) for target in batch_target]
from_state = torch.cat(from_state).to(device)
to_state = torch.cat(to_state).to(device)
trans_score = self.trans[from_state, to_state]
gather_target = torch.cat(batch_target).view(-1, 1).to(device)
emit_score = torch.gather(feats, 1, gather_target).to(device)
return trans_score.sum() + emit_score.sum()
def viterbi(self, feats, batch_input_lens):
word_size, tagset_size = feats.shape
batch_size = len(batch_input_lens)
viterbi_path = self.init_path(feats.shape) # use feats.shape to init path.shape
alpha = self.init_alpha(batch_size, tagset_size)
alpha[:, self.START_TAG_IDX, :] = 0 # Initialization
for legal_idx, legal_batch_size in self._iter_legal_batch(batch_input_lens):
feat = feats[legal_idx, :].view(legal_batch_size, 1, tagset_size)
legal_batch_score = feat + alpha[:legal_batch_size, ] + self.trans
alpha_new, best_tag = torch.max(legal_batch_score, 1).to(device)
alpha[:legal_batch_size, ] = alpha_new.unsqueeze(2)
viterbi_path[legal_idx, ] = best_tag
alpha = alpha + self.trans[:, self.END_TAG_IDX].unsqueeze(1)
path_score, best_tag = torch.max(alpha, 1).to(device)
path_score = path_score.squeeze() # path_score=#batch
best_paths = self.init_path((word_size, 1))
for legal_idx, legal_batch_size in self._iter_legal_batch(batch_input_lens, reverse=True):
best_paths[legal_idx, ] = best_tag[:legal_batch_size, ] #
backword_path = viterbi_path[legal_idx, ] # 1 * |Tag|
this_tag = best_tag[:legal_batch_size, ] # 1 * |legal_batch_size|
backword_tag = torch.gather(backword_path, 1, this_tag).to(device)
best_tag[:legal_batch_size, ] = backword_tag
# never computing <START>
# best_paths = #words
return path_score.view(-1), best_paths.view(-1)
class BiLSTM_CRF(nn.Module):
def __init__(self, vocab_size, tagset, embedding_dim, hidden_dim,
num_layers, bidirectional, dropout, start_tag, end_tag, device, pretrained=None):
super(BiLSTM_CRF, self).__init__()
self.bilstm = BiLSTM(vocab_size, tagset, embedding_dim, hidden_dim,
num_layers, bidirectional, dropout, pretrained)
self.CRF = CRF(tagset, start_tag, end_tag, device)
def init_hidden(self, batch_size, device):
self.bilstm.hidden = self.bilstm.init_hidden(batch_size, device)
def forward(self, batch_input, batch_input_lens, batch_mask):
feats = self.bilstm(batch_input, batch_input_lens, batch_mask)
score, path = self.CRF.viterbi(feats, batch_input_lens)
return path
def neg_log_likelihood(self, batch_input, batch_input_lens, batch_mask, batch_target):
feats = self.bilstm(batch_input, batch_input_lens, batch_mask)
gold_score = self.CRF.score_sentence(feats, batch_target)
forward_score = self.CRF.score_z(feats, batch_input_lens)
return forward_score - gold_score
def predict(self, batch_input, batch_input_lens, batch_mask):
return self(batch_input, batch_input_lens, batch_mask)
Training cell :
def prepare_sequence(seq, to_ix, device):
idxs = [to_ix[w] for w in seq]
return torch.tensor(idxs, dtype=torch.long).to(device)
def prepare_labels(lab, tag_to_ix, device):
idxs = [tag_to_ix[w] for w in lab]
return torch.tensor(idxs, dtype=torch.long).to(device)
class PadSequence:
def __call__(self, batch):
device = torch.device('cuda')
# Let's assume that each element in "batch" is a tuple (data, label).
# Sort the batch in the descending order
sorted_batch = sorted(batch, key=lambda x: len(x[0]), reverse=True)
# Get each sequence and pad it
sequences = [x[0] for x in sorted_batch]
sentence_in =[prepare_sequence(x, word_to_ix, device) for x in sequences]
sequences_padded = torch.nn.utils.rnn.pad_sequence(sentence_in, padding_value = len(word_to_ix) +1, batch_first=True).to(device)
lengths = torch.LongTensor([len(x) for x in sequences]).to(device)
masks = [True if index_word!=len(word_to_ix)+1 else False for sentence in sequences_padded for index_word in sentence ]
labels = [x[1] for x in sorted_batch]
labels_in = [prepare_sequence(x, tag_to_ix, device) for x in labels]
return sequences_padded, lengths, masks, labels_in
{ .... code to get the data formatted...}
device = torch.device("cuda")
batch_size = 64
START_TAG = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 200
HIDDEN_DIM = 20
NUM_LAYER = 3
BIDIRECTIONNAL = True
DROPOUT = 0.1
train_iter = DataLoader(dataset=training_data, collate_fn=PadSequence(), batch_size=64, shuffle=True)
model = BiLSTM_CRF(len(word_to_ix), tagset, EMBEDDING_DIM, HIDDEN_DIM, NUM_LAYER, BIDIRECTIONNAL, DROPOUT, START_TAG, STOP_TAG, device ).to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)
model.init_hidden(batch_size, device)
with tqdm(total=len(train_iter)) as progress_bar:
for batch_info in train_iter:
batch_input, batch_input_lens, batch_mask, batch_target = batch_info
loss_train = model.neg_log_likelihood(batch_input, batch_input_lens, batch_mask, batch_target)
optimizer.zero_grad()
loss_train.backward()
optimizer.step()
progress_bar.update(1) # update progress
Within PadSequence function (which acts as a collate_fn which gathers samples and makes a batch from them) you are explicitly casting to cuda device, namely:
class PadSequence:
def __call__(self, batch):
device = torch.device('cuda')
# Left rest of the code for brevity
...
lengths = torch.LongTensor([len(x) for x in sequences]).to(device)
...
return sequences_padded, lengths, masks, labels_in
You don’t need to cast your data when creating batch, we usually do that right before pushing the examples through neural network.
Also you should at least define the device like this:
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
or even better leave the choice of device for you/user in some part of the code where you setup everything.
The error lays in
lengths = torch.LongTensor([len(x) for x in sequences]).to(device)
change to
lengths = torch.LongTensor([len(x) for x in sequences]).cpu()