[Doc] Using train_fn with models where the attribute device has not been set

[nipunbatra]

Required prerequisites

• I have read this post from stackoverflow on creating a minimal, reproducible example
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version of ASTRA are you using?

0.0.2.dev21+g9573315

What Python version are you using?

3.9.7

Problem description

Using non-Atsra models with train_fn will cause the issue of device property not being set.

We can mention this upfront. In the error message also, we can tell what and how to set the property.

Also, in the Readme we can add some 1-2 lines about it.

Reproducible example code

The Python snippets:

import torch
import torch.nn as nn
from astra.torch.utils import train_fn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


### Vanilla Autoencoder 

class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(32 * 32 * 3, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.Linear(32, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 32 * 32 * 3),
            nn.Sigmoid()  # Sigmoid activation for pixel values between 0 and 1
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x
 model = Autoencoder().to(device)
 X = torch.randn(1000, 3, 32, 32).to(device)
 
criterion = nn.MSELoss()
iter_losses, epoch_losses = train_fn(model, X, X, criterion, lr=3e-4, 
                                     batch_size=128, epochs=30, verbose=False)
 

Error Traceback

File ~/miniforge3/lib/python3.9/site-packages/torch/nn/modules/module.py:1614, in Module.__getattr__(self, name)
   1612     if name in modules:
   1613         return modules[name]
-> 1614 raise AttributeError("'{}' object has no attribute '{}'".format(
   1615     type(self).__name__, name))

AttributeError: 'Autoencoder' object has no attribute 'device'

Additional context

No response

[patel-zeel]

Discussion

I think expecting models to have device property is too restrictive. All we want is to make sure both models and inputs are on the same device. More importantly, we want to support a case where model is on GPU but we can not put all data on GPU beforehand if it is too large, so, we want to pass to automatically infer the device from the model and put only a batch on the GPU.

Solution

I have removed device property from AstraModel and train_fn does not expect the models to have device property anymore. Instead, we have a function get_model_device in utils which gets the device of model from any torch model. train_fn uses get_model_device internally to infer model device.

Changes are implemented in 2dd2066