GPU Utilization Not Improving Training Time in PyTorch Regression Model

BeZeBeast: <p>I’ve implemented a simple linear regression model using PyTorch and attempted to accelerate training using a GPU. Despite this, I haven’t observed any improvement in training time. Profiling the model yields a significant CPU bottleneck rather than expected GPU optimization. Here’s a breakdown of the profiler’s output:</p>
<p>Top 10 Functions by CPU Time:<br>
model_fit: CPU time = 6268886.00 us<br>
Optimizer.step#FISTA.step: CPU time = 402481.00 us<br>
cudaLaunchKernel: CPU time = 302777.00 us<br>
aten::to: CPU time = 109666.00 us<br>
aten::<em>to_copy: CPU time = 107876.00 us<br>
aten::mul: CPU time = 85097.00 us<br>
aten::copy</em>: CPU time = 82672.00 us<br>
aten::abs: CPU time = 78300.00 us<br>
aten::quantile: CPU time = 75502.00 us<br>
aten::empty_strided: CPU time = 70781.00 us</p>
<p>Top 10 Functions by CUDA Time:<br>
model_fit: CUDA time = 41545.00 us<br>
Optimizer.step#FISTA.step: CUDA time = 23568.00 us<br>
aten::mul: CUDA time = 8151.00 us<br>
aten::sum: CUDA time = 5403.00 us<br>
aten::sub: CUDA time = 4679.00 us<br>
void at::native::vectorized_elementwise_kernel&lt;4, at::native::CUDAFunctor_add, at::detail::Array&lt;char*, 3&gt; &gt;(int, at::native::CUDAFunctor_add, at::detail::Array&lt;char*, 3&gt;): CUDA time = 4208.00 us<br>
aten::copy_: CUDA time = 4170.00 us<br>
aten::mse_loss: CUDA time = 3822.00 us<br>
aten::abs: CUDA time = 3601.00 us</p>
<p>It appears that a considerable amount of computation still happens on the CPU, especially multiplicative operations and data transfers (<code>aten::to</code> , <code>aten::_to_copy</code> ). Here’s the relevant part of my model code:</p>
<pre><code class="lang-auto">class LinearRegressionLoss(nn.Module):
    def __init__(self, lambda_, nu):
        super().__init__()
        self.lamb = lambda_
        self.nu = nu
        self.mse_loss = nn.MSELoss(reduction='sum')

    def forward(self, input, target, theta):
        square_root_lasso_loss = torch.sqrt(self.mse_loss(input, target))
        regularization = self.lamb * (torch.sum(torch.abs(theta.weight)) + torch.abs(theta.bias))
        total_loss = square_root_lasso_loss + regularization
        return total_loss, square_root_lasso_loss

class LinearModel():
    def __init__(self, lambda_):
        self.path_type = path_type
        self.lambda_= lambda_
        self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')

    def fit(self, X, y, verbose=False):
        # Data Processing
        X = torch.tensor(X.values, dtype=torch.float, device=self.device)
        scaler = StandardScaler()
        X = scaler.fit_transform(X)
        y = torch.tensor(y.values.squeeze(), dtype=torch.float, device=self.device)
        
        # Model Parameters
        self.theta = nn.Linear(X.shape[1], 1, dtype=torch.float, device=self.device)

        # Training
for i in range(-1, 6):
                init_lr = 0.1
                lambi = self.lambda_ * (np.exp(i) / (1 + np.exp(i)) if i &lt; 5 else 1)
                rel_err = 1e-11 if i==5 else 1e-5
                
                if verbose: 
print(f"Lambda = {lambi.item():.4f}")
                self.train_model(X, y, lambi, init_lr, rel_err, verbose) 
        if verbose: print("MODEL FITTED !")

    def forward(self, X):
        if X.device != self.device:
            X = X.to(self.device)

        return self.theta(X).squeeze()


    def train_model(self, X, y, lambda_, init_lr, rel_err, verbose):
        loss_fn = LinearRegressionLoss(lambda_, nu).to(self.device)
        train_score_fn = nn.MSELoss(reduction = 'mean').to(self.device)
        
        epoch, last_loss = 0, np.inf
        optimizer = FISTA(params=self.theta.parameters(), lr=init_lr, lambda_=lambda_)

        lr_factor = 0.9
        max_epochs = 10000

        while epoch &lt; max_epochs:
            optimizer.zero_grad()

            y_pred = self.forward(X)
            loss, bare_loss = loss_fn(y_pred, y, self.theta)
            loss = loss.detach()
            train_loss = train_score_fn(y_pred, y).detach()

            if loss &gt; last_loss: 
                learning_rate = optimizer.param_groups[0]['lr']
                optimizer = FISTA(params=self.theta.parameters(), lr=learning_rate*lr_factor, lambda_=lambda_)

            if epoch % 20 == 0:
                if verbose:
                    print(f"\tEpoch: {epoch} | Loss: {loss.item():.5f}")

                if epoch &gt; 0 and abs(loss - last_loss) / loss &lt; rel_err:
                    if verbose: print(f"\n\t Descent stopped: loss is no longer decreasing significantly.\n")
                    break
                
                last_loss = loss

            epoch += 1
            bare_loss.backward()
            optimizer.step()

            if epoch == max_epochs and verbose: print("FISTA descent stopped: maximum iterations reached") 
</code></pre>
<p>And the optimizer:</p>
<pre><code class="lang-auto">class FISTA(torch.optim.Optimizer):
    def __init__(self, params, lr, lambda_):
        self.lr = lr
        self.lambda_ = lambda_
        defaults = dict(lr=lr)
        super(FISTA, self).__init__(params, defaults)
        

    def shrinkage_operator(self, u, lambda_):
        '''Applies the shrinkage operator to a PyTorch tensor.'''
            return u.sign() * torch.clamp(u.abs() - lambda_, min=0.0)
        
    @torch.no_grad()
    def step(self, closure=None):
        '''Performs a single optimization step.'''
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad
                state = self.state[p]

                if 'x_prev' not in state:
                    # We use .detach() to ensure we do not track history
                    state['x_prev'] = p.detach().clone()
                    state['y_prev'] = p.detach().clone()
                    state['t_prev'] = torch.tensor(1., device=p.device)

                x_prev, y_prev, t_prev = state['x_prev'], state['y_prev'], state['t_prev']

                x_next = self.shrinkage_operator(y_prev - self.lr * grad, self.lr * self.lambda_)
                t_next = (1. + torch.sqrt(1. + 4. * t_prev ** 2)) / 2.
                y_next = x_next + ((t_prev - 1) / t_next) * (x_next - x_prev)

                state['x_prev'].copy_(x_next)
                state['y_prev'].copy_(y_next)
                state['t_prev'].copy_(t_next)

                p.copy_(x_next)

        return loss
</code></pre>
<ol>
<li>Why is the GPU underutilized, and why are operations like <code>aten::mul</code> and <code>aten::abs</code> still performed on the CPU?</li>
<li>How can I ensure that more computation, especially tensor operations, is offloaded to the GPU to improve training efficiency?</li>
</ol>
<p>Any insights or suggestions on how to better leverage GPU capabilities in this PyTorch model would be greatly appreciated!</p>
<hr>
<p>Edit : I added a ‘device’ argument to the model class such that if given, the model’s device is forced. Results:</p>
<pre><code class="lang-auto">t = time.time()
model = LinearModel(path_type=0)
model.fit(X, y, False)
print(time.time()-t)

--&gt;0.7786831855773926

t = time.time()
model = LinearModel(path_type=0, device='cpu')
model.fit(X, y, False)
print(time.time()-t)

--&gt; 0.800177001953125
</code></pre>