How to Load a 'Learner' Model with a Custom Loss Function in a Flask Application

Deep Learning
1

Hi all
I am currently working on loading a learner model from a pickle file. This model includes a custom loss function and needs to be integrated into a Flask application. However, I keep encountering the following error:

Custom classes or functions exported with your `Learner` not available in namespace.\Re-declare/import before loading:
        Can't get attribute 'combined_loss' on <module '__main__' from 'C:\\Users\\Desktop\\Meteor\\flask_app\\venv\\Scripts\\flask.exe\\__main__.py'>

Does anyone know how I can resolve this issue?

1 Like
2

I think you need to add the code for your custom loss function to your flask app. They are not exported automatically

1 Like
3

Thanks @Archaeologist,
I have tried but I still get the same error

4

If you like, please show your code here

1 Like
5

@Archaeologist Thanks again!!

This is the main script:

import io
import torch
import pathlib
import numpy as np
from utils import *
import torch.nn.functional as F

from fastai.vision.all import * 
from flask import Flask, request, jsonify, send_file, render_template

app = Flask(__name__)

MODEL_PATH = "C:\\Users\\abc\\Desktop\\xyz\\Meteor\\flask_app\\export_large.pkl"
IMG_SIZE = 512

# Custom loss
def regularization_loss(y_pred):
    dx = torch.abs(y_pred[:, :, 1:] - y_pred[:, :, :-1])
    dy = torch.abs(y_pred[:, :, :, 1:] - y_pred[:, :, :, :-1])
    reg_loss = torch.mean(dx) + torch.mean(dy)
    return reg_loss

def combined_loss(y_pred, y_true):
    ce_loss = F.cross_entropy(y_pred, y_true, weight=wgts)
    reg_loss = regularization_loss(y_pred)
    total_loss = ce_loss + reg_loss 
    return total_loss

# Load model
posix_backup = pathlib.PosixPath 
pathlib.PosixPath = pathlib.WindowsPath 
learn = load_learner(MODEL_PATH, cpu=True)
pathlib.PosixPath = posix_backup 
learn.dls.to(device='cpu') # device='cuda'
learn.model.to(device='cpu') # device='cuda'

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({"error": "Select a file!!"}), 400
    file = request.files['file']
    try:
        img_resized_8bit = from_16_to_8(file, IMG_SIZE) 
    except Exception as e:
        return jsonify({"error": str(e)}), 400

    pred_mask, _, _ = learn.predict(img_resized_8bit)
    pred_mask_np = np.array(pred_mask[0].numpy())
    threshold = 0.5
    binary_mask = (pred_mask_np > threshold).astype(np.uint8) 
    binary_mask_image = Image.fromarray(binary_mask * 255)

    y_coords, x_coords  = find_highest_lowest_xy(binary_mask_image)

    buf = plot_coords(img_resized_8bit, x_coords, y_coords)

    return send_file(buf, mimetype='image/png')

if __name__ == '__main__':
    app.run(debug=True)
6

Seems OK.

The only thing I noticed is that you are not using any Fastai loss but pure pytorch losses.

Note that they are not fully compatible. You might want to look at this conversation for example :

1 Like
7

I just can’t understand why it doesn’t work

I tried exporting the model in .dill format, it works now, but unfortunately the model’s performance has deteriorated a lot