LSTM on GPU still working on CPU

alexmulo: <p>Dear all,</p>
<p>I did some research but I could not find anything relevant.<br>
I have LSTM and I would like it to work completely on the GPU to speed up training.<br>
Everything works fine but slow because the system is partially using the GPU. I trained the same data set with similar LSTM model on keras and it is way faster. Do you have any idea what is wrong with my code?</p>
<p>I have the latest version of Pytorch and I am using Ubuntu 20.04 with an NVIDIA GTX 1070 with CUDA 11.2 .</p>
<p>This is my code:</p>
<pre><code class="lang-auto">class LSTM_CUDA(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
        super(LSTM_CUDA, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers,batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(1), self.hidden_dim).cuda()
        c0 = torch.zeros(self.num_layers, x.size(1), self.hidden_dim).cuda()
        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))
        out = self.fc(out[-1])
        return out

model_cuda = LSTM_CUDA(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers)
model_cuda = model_cuda.cuda()
criterion = torch.nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimiser = torch.optim.Adam(model_cuda.parameters(), lr=0.01)

def validation_metrics (model, valid_dl):
    model.eval()
    correct = 0
    total = 0
    sum_loss = 0.0
    sum_rmse = 0.0
    for batch_idx, (x, y) in enumerate(valid_dl):
        x = x.cuda()
        y = y.long().cuda()
        y_hat = model(x)
        loss = criterion(y_hat, y)
        pred = torch.max(y_hat, 1)[1]
        if torch.cuda.is_available():
            correct += (pred.cpu() == y.cpu()).sum()
        else:
            correct += (pred == y).sum()
        #correct += (pred == y).float().sum()
        total += y.shape[0]
        sum_loss += loss.item()*y.shape[0]
        if torch.cuda.is_available():
            y = y.cpu()
            pred = pred.cpu()
            
        sum_rmse += np.sqrt(mean_squared_error(pred, y.unsqueeze(-1)))*y.shape[0]
    return sum_loss/total, correct/total, sum_rmse/total

import time
start_time = time.time()

hist = np.zeros(10)
sum_loss = 0.0
total = 0
for t in range(num_epochs):
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        model_cuda.train()
        optimiser.zero_grad()
        x = inputs.cuda()
        y = targets.long().cuda()
        y_train_pred = model_cuda(x)    

        loss = criterion(y_train_pred, y)
        sum_loss += loss.item()*y.shape[0]
        total += y.shape[0]
        loss.backward()
        optimiser.step()
    print("Epoch ", t, "Loss train: ", loss.item())
    hist[t] = loss.item()
    val_loss, val_acc, val_rmse = validation_metrics(model_cuda, val_tensor)
    print("train loss %.3f, val loss %.3f, val accuracy %.3f, and val rmse %.3f" % (sum_loss/total, val_loss, val_acc, val_rmse))
    
print("--- %s seconds ---" % (time.time() - start_time))
</code></pre>
<p>Thanks.</p>

J_Johnson: <p>I usually just use the “to” method:</p>
<pre><code class="lang-auto">device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Assuming that we are on a CUDA machine, this should print a CUDA device:

print(device)

model_cuda.to(device) # This replaces: model_cuda = model_cuda.cuda()
</code></pre>
<p>Please see here for more information:<br>
<a href="https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#training-on-gpu" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#training-on-gpu</a></p>

alexmulo: <p>Hi,</p>
<p>thanks for your reply. Unfortunately nothing really changed. Same slow speed of ca. 54 seconds.</p>
<p>Any other suggestions?</p>
<p>Thanks.</p>

J_Johnson: <p>What about this:</p>
<blockquote>
<p>WARNING</p>
<p>There are known non-determinism issues for RNN functions on some versions of cuDNN and CUDA. You can enforce deterministic behavior by setting the following environment variables:</p>
<p>On CUDA 10.1, set environment variable <code>CUDA_LAUNCH_BLOCKING=1</code>. This may affect performance.</p>
<p>On CUDA 10.2 or later, set environment variable (note the leading colon symbol) <code>CUBLAS_WORKSPACE_CONFIG=:16:8</code> or <code>CUBLAS_WORKSPACE_CONFIG=:4096:2</code>.</p>
<p>See the <a href="https://docs.nvidia.com/deeplearning/sdk/cudnn-release-notes/rel_8.html" rel="noopener nofollow ugc">cuDNN 8 Release Notes</a> for more information.</p>
</blockquote>
<p>Ref: <a href="https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html#torch.nn.LSTM" class="inline-onebox" rel="noopener nofollow ugc">LSTM — PyTorch 2.1 documentation</a></p>