Cnn regression handling

vision
1

Hello !

I want to adapt a model from segmentation_model_pytorch to predict a continuoiuys value, regression.
So i have image and a label (in range 0 to 100)

what should i do ? i am a beginner and i did a dummy model:

class CustomModel(nn.Module):
    def __init__(self, cfg, weight=None):
        super().__init__()
        self.cfg = cfg

        self.encoder = smp.Unet(
            encoder_name=cfg.backbone, 
            encoder_weights=weight,
            classes=cfg.target_size,
            activation='identity',
        )
        self.global_pool = nn.AdaptiveAvgPool2d(1)
    def forward(self, image):
        output = self.encoder(image)
        output = self.global_pool(output)
        output = output.view(output.size(0), -1)  # Flatten
        return output

def build_model(cfg, weight="imagenet"):
    print('model_name', cfg.model_name)
    print('backbone', cfg.backbone)

    model = CustomModel(cfg, weight)
    return model

and i try a random tensor :

import torch
import torch.nn.functional as F
input_tensor = torch.randn(1, 3, 256, 256)
weight = "imagenet"  
model = build_model(CFG, weight)
output = model(input_tensor)
regression_output = output.squeeze().cpu().detach().numpy()
print(regression_output)

Is it all i should do ?? i mean just set the activation to none (plus avg pool/flatten?) shouldn’t i remove the segmentation head and the decoder part?
I am a bit lost, if someone can just guide me ? It seems logical that i just need the encoder part right?
Have a good day !

2

Sorry i think because its more in the ecosystem but not pytorch i dont get any response, or maybe my question does not make any sense!

But the question is more:
Using a segmentation library and modifying only the lasty activation function to be linaer is ok for a regression task (i am pretty sure yes because it seems working on my side) ?
Or i should remove the segmentation head and the decoder part to get better result?

Have a good day!

3

wow i am worst than i already thought, i was already doing only on the encoder …
well i am pretty sure is ok:

class CustomModel(nn.Module):
    def __init__(self, cfg, weight=None):
        super().__init__()
        self.cfg = cfg

        self.encoder = smp.Unet(
            encoder_name=cfg.backbone, 
            encoder_weights=weight,
            classes=cfg.target_size,
            activation='identity',
        )
        self.fc = nn.LazyLinear( out_features=1) 

    def forward(self, image):
        output = self.encoder(image)
        output = output.view(output.size(0), -1) 
        output = self.fc(output)  
        return output