GPU Running for Pyro using MyModel().to(device) not responding

mmabaei: <p>Dear All,</p>
<p>I am new to using PyTorch and Pyro with a GPU. I already have my model, but it takes a long time to run on a CPU. That’s why I’m considering shifting to GPU. However, I’m not sure if my system is running on the GPU or not, even though <code>torch.device</code> is successfully set to CUDA. I’m wondering if the <code>to.(device)</code> command will automatically instruct Pyro to run on the GPU.</p>
<p>When I check my system performance during runtime, the GPU usage is almost 0%, but the CPU usage varies. I’ve tried using a simple code to test GPU usage. I would greatly appreciate any advice or help you can provide.</p>
<p>Thank you.<br>
import torch<br>
import pyro<br>
from pyro.nn import PyroModule, PyroSample<br>
import pyro.distributions as dist<br>
from pyro.infer import MCMC, NUTS</p>
<h1><a name="check-if-cuda-is-available-and-set-the-device-accordingly-1" class="anchor" href="#check-if-cuda-is-available-and-set-the-device-accordingly-1"></a>Check if CUDA is available and set the device accordingly</h1>
<p>device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
print(f"Running on device: {device}")</p>
<h1><a name="define-a-simple-pyro-model-using-pyromodule-2" class="anchor" href="#define-a-simple-pyro-model-using-pyromodule-2"></a>Define a simple Pyro model using pyro.Module</h1>
<p>class MyModel(PyroModule):<br>
def <strong>init</strong>(self):<br>
super().<strong>init</strong>()</p>
<pre><code>def forward(self):
    x = pyro.sample("x", dist.Normal(0,   1))
    y = pyro.sample("y", dist.Normal(x,   1))
    return y
</code></pre>
<h1><a name="instantiate-the-model-and-move-it-to-the-gpu-3" class="anchor" href="#instantiate-the-model-and-move-it-to-the-gpu-3"></a>Instantiate the model and move it to the GPU</h1>
<p>model = MyModel().to(device)</p>
<h1><a name="perform-inference-using-mcmc-4" class="anchor" href="#perform-inference-using-mcmc-4"></a>Perform inference using MCMC</h1>
<p>nuts_kernel = NUTS(model.forward)<br>
mcmc = MCMC(nuts_kernel, num_samples=100000, warmup_steps=500)<br>
mcmc.run()</p>

ptrblck: <p>Could you try to run e.g. the  <a href="https://github.com/pytorch/examples/tree/main/mnist"><code>MNIST</code> example</a> and check if your GPU is used?</p>

mmabaei: <p>Thank you for your valuable advice. I implemented the MNIST example and closely monitored the performance of my NVIDIA GPU. Surprisingly, the GPU appears to contribute only 3% to the overall computation, while the majority of simulations run on the CPU. I’m uncertain whether this distribution is reasonable or if there might be an issue with the GPU connection? Should I do further amendment with my installation? Below, you’ll find the performance plot for reference:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/b/5bbb446ffdac68c098e7252d5b4ad02f62f78623.png" data-download-href="https://discuss.pytorch.org/uploads/default/5bbb446ffdac68c098e7252d5b4ad02f62f78623" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/b/5bbb446ffdac68c098e7252d5b4ad02f62f78623.png" alt="image" data-base62-sha1="d5uFqEiWYOBKLKEVetm0u6aorWb" width="690" height="481" data-dominant-color="F5F7F7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1053×735 30.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>The Windows Task Manager might be misleading as you would need to enable the <code>compute</code> tab as also described <a href="https://discuss.pytorch.org/t/it-seems-pytorch-doesnt-use-gpu/74673/6">here</a>.<br>
Alternatively, use <code>nvidia-smi</code> as it should show the utilization.<br>
However, assuming you’ve moved the data and model to the GPU it seems to work as you don’t see a hang, so I assume the original issue is unrelated to the GPU usage.</p>

mmabaei: <p>Thank you so much again for your advice which was very helpful. IAfter Minst data set, I tried to run a Deep Learning model consist of 70 leyers of LSTM each of with 100 hidden size and the NVIDIA GPU contributed for 70% of computaion in less than a minute!.   Here is the plot:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/6/a68321cb2b348b19c9c5edc38a094c9a70c34dea.png" data-download-href="https://discuss.pytorch.org/uploads/default/a68321cb2b348b19c9c5edc38a094c9a70c34dea" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a68321cb2b348b19c9c5edc38a094c9a70c34dea_2_690x306.png" alt="image" data-base62-sha1="nL2cXRJvAT5qx5zNZwsSmOps6bU" width="690" height="306" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a68321cb2b348b19c9c5edc38a094c9a70c34dea_2_690x306.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a68321cb2b348b19c9c5edc38a094c9a70c34dea_2_1035x459.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a68321cb2b348b19c9c5edc38a094c9a70c34dea_2_1380x612.png 2x" data-dominant-color="F7F9FA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1514×672 30.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I have a Pyro model in which i use probablistic terms for  LSTM blockes, when I do the same process to transfer the model, parameter and data into deveice, the python give bellow error:</p>
<p>RuntimeError: Input and parameter tensors are not at the same device, found input tensor at cuda:0 and parameter tensor at cpu</p>
<p>however I have the same model just transfer to pyro for Bayesian Deep Learning. this my Pyro model:</p>
<p>class Model(PyroModule):<br>
def <strong>init</strong>(self, input_size=1, num_classes=1, hidden_size=3, num_layers=1, prior_scale=50.0):<br>
super(Model,self).<strong>init</strong>()</p>
<pre><code>    self.num_classes = num_classes
    self.num_layers = num_layers
    self.input_size = input_size
    self.hidden_size = hidden_size
    
    self.activation = nn.ReLU()  # or nn.ReLU()
    # Correctly initialize the LSTM layer with bidirectional=False
    self.lstm = PyroModule[nn.LSTM](input_size, hidden_size, num_layers, batch_first=True, bidirectional=False)
    self.linear = PyroModule[nn.Linear](hidden_size,  128)  # Adjusted for unidirectional LSTM
    self.fc = PyroModule[nn.Linear](128, num_classes)

    # Initialize weights and biases for each layer
    # Input to hidden layer
    self.lstm.weight_ih_l0 = PyroSample(dist.Normal(0., prior_scale).expand([4*hidden_size, input_size]).to_event(2))
    self.lstm.bias_ih_l0 = PyroSample(dist.Normal(0., prior_scale).expand([4*hidden_size]).to_event(1))
    
    # Hidden to hidden layer
    self.lstm.weight_hh_l0 = PyroSample(dist.Normal(0., prior_scale).expand([4*hidden_size, hidden_size]).to_event(2))
    self.lstm.bias_hh_l0= PyroSample(dist.Normal(0., prior_scale).expand([4*hidden_size]).to_event(1))

    self.linear.weight = PyroSample(dist.Normal(0., prior_scale).expand([128,hidden_size]).to_event(2))
    self.linear.bias = PyroSample(dist.Normal(0., prior_scale).expand([128]).to_event(1))
    self.fc.weight = PyroSample(dist.Normal(0., prior_scale).expand([num_classes, 128]).to_event(2))
    self.fc.bias = PyroSample(dist.Normal(0., prior_scale).expand([num_classes]).to_event(1))
    
def forward(self, x, y=None,noise_shape = 0.5):

    # Initialize hidden state and cell state
    h_0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) #hidden state
    c_0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) #internal state
    
    
    output, (hn, cn) = self.lstm(x, (h_0, c_0)) #lstm with input, hidden, and internal state
         
    hn = hn.view(-1, self.hidden_size) #reshaping the data for Dense layer next
    
    out = self.activation(hn)
    out = self.linear(out)
    out = self.activation(out)
    mu = self.fc(out)
    sigma = pyro.sample("sigma", dist.Gamma(noise_shape, 1))  # infer the response noise

    with pyro.plate("data", x.shape[0]):
        obs = pyro.sample("obs", dist.Normal(mu, sigma * sigma), obs=y)
    return mu
</code></pre>
<p>and this my setup for GPU running:</p>
<h1><a name="check-for-gpu-availability-1" class="anchor" href="#check-for-gpu-availability-1"></a>Check for GPU availability</h1>
<p>use_cuda = torch.cuda.is_available()<br>
device = torch.device(“cuda” if use_cuda else “cpu”)</p>
<h1><a name="initialize-the-bayesian-lstm-model-2" class="anchor" href="#initialize-the-bayesian-lstm-model-2"></a>Initialize the Bayesian LSTM model</h1>
<p>input_size = X_test_tensors_final.shape[2] # Number of features<br>
hidden_size = 2 # Number of features in hidden state<br>
num_layers = 1 # Number of stacked LSTM layers<br>
model = Model(input_size=input_size, num_classes=1, hidden_size=hidden_size,<br>
num_layers=num_layers, prior_scale=50.0).to(device)</p>
<h1><a name="set-up-the-mcmc-sampler-3" class="anchor" href="#set-up-the-mcmc-sampler-3"></a>Set up the MCMC sampler</h1>
<p>nuts_kernel = NUTS(model)<br>
mcmc = MCMC(nuts_kernel, num_samples=50, warmup_steps=50)</p>
<h1><a name="create-a-combined-dataset-4" class="anchor" href="#create-a-combined-dataset-4"></a>Create a combined dataset</h1>
<p>dataset = TensorDataset(X_train_tensors_final, y_train_tensors)</p>
<h1><a name="set-the-batch-size-5" class="anchor" href="#set-the-batch-size-5"></a>Set the batch size</h1>
<p>batch_size = 400</p>
<h1><a name="create-a-dataloader-6" class="anchor" href="#create-a-dataloader-6"></a>Create a DataLoader</h1>
<p>data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)</p>
<h1><a name="initialize-an-empty-list-to-store-mcmc-results-7" class="anchor" href="#initialize-an-empty-list-to-store-mcmc-results-7"></a>Initialize an empty list to store MCMC results</h1>
<p>all_mcmc_results = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<h1><a name="iterate-over-the-batches-8" class="anchor" href="#iterate-over-the-batches-8"></a>Iterate over the batches</h1>
<p>for batch_idx, (X_batch, y_batch) in enumerate(data_loader):<br>
# Move data to GPU if available<br>
X_batch = X_batch.to(device)<br>
y_batch = y_batch.to(device)</p>
<pre><code># Run your MCMC simulation
 mcmc_result = mcmc.run(X_batch, y_batch)

# Append the result to the list
all_mcmc_results.append(mcmc_result)
</code></pre>

ptrblck: <p>I don’t know how Pyro interacts with your model and you might need to check the stacktrace trying to isolate the failing line of code.</p>