PyTorch Model Demo - Load the model & Inference only ONCE

[TdProData]

I have developped a streamlit app which loads a pretrained custom model from PyTorch, with a selectbox widget which enables the user to select the name of the frame from the test set to use for inference. Then, this input is sent to the model whose output is also displayed.

Here is the problem : as streamlit always reads code in a loop, the application loads the model and repeats the same inferences in a loop at each iteration. Is it possible to force it not to load the model in a loop?

In other words, I want to load the model once and then have the freedom to display the frame I want and perform inference, without the model reloading again.

Code snippet:

# Load model
embed_size = st.selectbox("Which emb. dim. ?", ("256", "512"))
loss_type = st.selectbox("Which loss ?", ("L1", "L1"))
model = loadPretrainedModel(embed_size, loss_type)

# Load names of test frames
frames_name = getTestNameFrames() # list with name of frame
selected_img = st.selectbox("Which frame ?", frames_name)
img = getImage(selected_img) # load selected frame as np array
st.image(img, caption="Selected frame")

# perform inference and display output
output = model(img)
st.image(output, caption="Output")

Thank you for your help !

[DirkRemmers]

Hey TdProData,

I think you might be fine with caching the function where you load up the model.
If you use:

@st.cache_resource()
def load_model(embed_size, loss_type):
    return loadPretrainedModel(embed_size, loss_type)

the loading of the model is cached, and will be much faster in the next rerun.
If you want to reload the model using that aproach, just call load_model.clear() and rerun the page.

Another approach you could use, is to put the model inside session state:

if "model" not in st.session_state.keys():
    st.session_state["model"] = loadPretrainedModel(embed_size, loss_type)
model = st.session_state["model"]

If you want to reload the model using this approach, just call del st.session_state["model"] and rerun the page.

Hopefully this solves your problem!

Cheers,
Dirk

[TdProData]

Way better now, thanks !