I’m training a language model in colab on GPU but it’s really slow, How can I optimize my code?
This is my code:
- The training:
for epoch in range(epochs):
states = (torch.FloatTensor(2, bs, 200).\
uniform_(-0.1, 0.1).to(device),\
torch.FloatTensor(2, bs, 200).\
uniform_(-0.1, 0.1).to(device))
states = detach(states)
for i, batch in enumerate(train_iter):
model.train()
inputs, labels = batch.text, batch.target
inputs, labels = inputs.to(device), labels.to(device)
outputs, states = model(inputs, states)
loss = loss_fn(outputs, labels.view(-1))
model.zero_grad()
loss.backward(retain_graph=True)
clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
if not i%100:
print('i {}, Epoch [{}/{}], Loss: {:.4f}, Perplexity: {:5.2f}'
.format(i, epoch+1, epochs, loss.item(), np.exp(loss.item())))
- The model:
class RNNModel(nn.Module):
def __init__(self,
rnn_type,
vocab_size=10001,
n_layers=2,
n_hidden=200,
emb_size=300,
p_dropout=0.,
batch_size=bs):
super(RNNModel, self).__init__()
self.embedding = nn.Embedding.from_pretrained(TEXT.vocab.vectors, freeze=True)#(vocab_size, emb_size)
#self.embedding.weight.data.copy_(TEXT.vocab.vectors)
self.rnn = getattr(nn, rnn_type)(input_size=emb_size,
num_layers=n_layers,
hidden_size=n_hidden,
dropout=p_dropout,
batch_first=True)
self.linear = nn.Linear(n_hidden,
vocab_size)
def forward(self, x, h):
x = self.embedding(x)
output, (h, c) = self.rnn(x, h)
output = output.reshape(output.size(0)*output.size(1),
output.size(2))
output = self.linear(output)
return output, (h, c)