<jupyter_start><jupyter_text>The validity of JIT model<jupyter_code>from pathlib import Path
import os
HOME = Path(os.environ['HOME'])
MODELNAME = "fpn"
ENCODER = "efficientnet-b5"
DOWN = False # Downsample at the bottom
SRC = HOME/"ucsi"/"fastai"/"models"/"bestmodel_3.pth" # source model path
DST = HOME/"ucsi"/"jit"/"fpn_b5_e3.pth" # desitination model path
from torch import jit
import segmentation_models_pytorch as smp
import torch
from torch import nn
if MODELNAME =="fpn":
    model_class = smp.FPN
elif MODELNAME == "unet":
    model_class = smp.Unet<jupyter_output><empty_output><jupyter_text>Loading The Model<jupyter_code>seg_conf = {
    "encoder_name":ENCODER,
    "encoder_weights":None,
    "classes":4,
    "activation":"sigmoid",
}

print("Constructing the model")
print(seg_conf)
if DOWN:
    class majorModel(nn.Module):
        def __init__(self, seg_model):
            super().__init__()
            self.seq = nn.Sequential(*[
                nn.Conv2d(3,12,kernel_size=(3,3), padding=1, stride=1, ),
                nn.ReLU(),
                nn.Conv2d(12,3,kernel_size=(3,3), padding=1, stride=2),
                nn.ReLU(),
                seg_model,])
            
        def forward(self,x):
            return self.seq(x)
    model = majorModel(model_class(**seg_conf))
    
else:
    model = model_class(**seg_conf)
CUDA = torch.cuda.is_available()
print("CUDA available:\t%s"%(CUDA))
print("Loading from weights:\t%s"%(SRC))
state = torch.load(SRC)
if "model" in state:
    state = state["model"]
if CUDA:
    model = model.cuda()
model.load_state_dict(state)
testimg = torch.rand(2, 3, 320, 640)
if CUDA:
    testimg = testimg.cuda()
model = model.eval()
with torch.no_grad():
    y1 = model(testimg)<jupyter_output><empty_output><jupyter_text>Save to JIT<jupyter_code>print("Saving to jit traced model:\t%s"%(DST))
with torch.no_grad():
    traced = jit.trace(model, testimg)
    traced.save(str(DST))
if CUDA:
    model = model.cpu()<jupyter_output><empty_output><jupyter_text>Recover from saved JIT<jupyter_code>recovered = jit.load(str(DST))
if CUDA: 
    recovered = recovered.cuda()
with torch.no_grad():
    y2 = recovered(testimg)
print("Absolute Mean Error:%s"%(torch.abs(y1-y2).mean().item()))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14689016321129245, 'fraction_numerical': 0.014556682840758712, 'mean_word_length': 3.9197396963123645, 'pattern_counts': {'":': 6, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '179178', 'license': 'MIT', 'n_tokens_mistral': 864, 'n_tokens_neox': 835, 'n_words': 160, 'path': 'le_jit.ipynb', 'repo_name': 'iofthetiger/ucsi', 'size': '2267'}