getting raw logits from trained model

[snachi2s]

问题确认 Search before asking

• [X] 我已经搜索过问题，但是没有找到解答。I have searched the question and found no related answer.

请提出你的问题 Please ask your question

Hi, I trained Unet++ model with my custom dataset and I want to do some post-processing over the predicted masks. Is there any way to get the predictions (in the form of logits or arrays) during inference?

Following is the code I used for API inferencing,

from paddleseg.models import unet_plusplus
import paddle 
import paddleseg.transforms as T

model = unet_plusplus.UNetPlusPlus(
    num_classes=4, 
    in_channels=3, 
    use_deconv=False, 
    pretrained=None,
    align_corners=False,
    is_ds=True)

model_path = 'trained_logs/unet++/best_model/model.pdparams'
image_list = ['data/image.png']


if model_path:
    para_state_dict = paddle.load(model_path)  
    model.set_dict(para_state_dict)            # Load parameters
    print('Loaded trained params of model successfully')
else:
    raise ValueError('The model_path is wrong: {}'.format(model_path))

transforms = T.Compose([
    T.Resize(target_size=(256,256)),
    T.Normalize()
])


from paddleseg.core import predict
predict(
    model,
    model_path=model_path,
    transforms=transforms,
    image_list=image_list,
    save_dir='best_unet++_model',
)

Thanks!!