# Regression(MSELoss())/images/Inference

**I am using MSE loss() for a regression problem but with images. I want to be able to predict the outcome from a test set of images using my saved model. (Given below is a sample code i used but didnt work well). Also i am unsure if it is the right method to go about it.
How can i do that? (Sample code would be helpful)

During training my my training and validation loss doesnt improve much. Thanks**

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

model = UNet1()

model.eval()

image = np.expand_dims(image, axis=0)

image = np.expand_dims(image, axis=1)

images = torch.from_numpy(image).float()

#Predictions

fig = plt.figure(figsize=(20, 12))

rows = 1

columns = 3

# Adds a subplot at the 1st position

plt.imshow(images.squeeze(0).permute(1,2,0).squeeze(2).numpy(), cmap=‘gray’)

plt.axis(‘off’)

plt.title(“Image”)

# Adds a subplot at the 2nd position

plt.axis(‘off’)

plt.title(“Prediction”)

# Adds a subplot at the 3rd position

plt.axis(‘off’)

I would recommend to scale down the use case and try to overfit a small dataset (e.g. just 10 samples) by playing around with some hyperparameters and make sure your model is able to learn them. Once this is done you could scale up the use case again.

PS: you can post code snippets by wrapping them into three backticks ```, which makes debugging easier.

1 Like

Hey thanks for getting back. If I am doing inference on an image for a regression problem, is it similar to doing with a classification problem or just pass the image though the model and get the outputs.
I am using the mse loss function

Ex: train a classifier example from pytorch on cifar

Yes, you can pass the input to the model and use the output as the regression prediction. Since you would usually not update any parameters during evaluation, you could also wrap the forward pass into a `with torch.no_grad()` guard to save memory.

1 Like

For inference:

I tried torch.argmax(model(outputs), dim=1) with the torch.no grad wrapper. I plotted the outcome and got zeros

Model(outputs) has values when I plot them though without argrmax function

A regression task usually uses floating point outputs, while `torch.argmax` will return the index in the specified dimension containing the max. value and is thus commonly used to get the class predictions for a multi-class classification use case.

What should I use then for regression? An example would help

You can use the output of the last layer directly:

``````model = nn.Sequential(
nn.Linear(10, 10),
nn.ReLU(),
nn.Linear(10, 10)
)
x = torch.randn(1, 10)
regression_output = model(x)
``````
1 Like

Thanks for the immediate response and help.