[{"data":1,"prerenderedAt":1721},["ShallowReactive",2],{"blog-\u002Fblog\u002F2025\u002F10\u002Fcustom-onnx-model":3},{"id":4,"title":5,"body":6,"description":12,"extension":1711,"meta":1712,"navigation":807,"path":1717,"seo":1718,"stem":1719,"__hash__":1720},"blog\u002Fblog\u002F2025\u002F10\u002Fcustom-onnx-model.md","Deploy Custom-Trained AI Models: Using ONNX with Node-RED and FlowFuse",{"type":7,"value":8,"toc":1701},"minimark",[9,13,18,21,32,36,39,43,46,59,62,66,71,79,85,90,100,132,141,145,156,205,209,212,226,229,246,249,263,269,272,314,317,369,376,395,399,413,638,645,757,761,764,770,788,791,1439,1443,1446,1457,1460,1512,1519,1523,1602,1611,1619,1623,1627,1630,1642,1649,1680,1683,1697],[10,11,12],"p",{},"FlowFuse is introducing a new set of AI nodes to make it easier than ever to integrate AI and machine learning into your Node-RED workflows.\nIn this guide, you will learn how to train an image classifier model, and use it with the new FlowFuse AI Nodes to recognise your own products, components - or anything else you can imagine.",[14,15,17],"h3",{"id":16},"introduction","Introduction",[10,19,20],{},"In this article, we will be building a PyTorch-based image classification model to identify fruit types (apple, kiwi, mango) using a dataset of labelled images.\nOf course, you would typically be classifying your own things like your company widgets and products, but for the sake of learning the process, we will be using images of fruit.\nOnce the model is trained, it is exported to the ONNX format, it is then ready for use with the new FlowFuse AI nodes.",[10,22,23,24,31],{},"Note: The code and sample dataset used in this tutorial can be downloaded from ",[25,26,30],"a",{"href":27,"rel":28},"https:\u002F\u002Fwebsite-data.s3.eu-west-1.amazonaws.com\u002F2025-10-onnx-model-training-dataset.zip",[29],"nofollow","this link",".",[14,33,35],{"id":34},"some-background-first","Some background first",[10,37,38],{},"The process we will use is commonly referred to as \"transfer learning\". This is where you take a pre-trained model and fine-tune it on your own dataset.\nThis is a common approach in deep learning as it allows us to leverage the knowledge learned by the pre-trained model and adapt it to our specific task with a smaller dataset.  For reference, this tutorial will use ResNet-18 which is an 18-layer Residual Network (ResNet), a convolutional neural network (CNN) architecture that uses \"skip connections\" to help train very deep networks by addressing the vanishing gradient problem. Pre-trained ResNet-18 models are often trained on the ImageNet dataset and are widely used for image classification of 1000 categories.",[14,40,42],{"id":41},"overview-of-operations","Overview of operations",[10,44,45],{},"The 3 main steps to achieve this involves:",[47,48,49,53,56],"ol",{},[50,51,52],"li",{},"Setting up a Python environment with PyTorch, TorchVision, ONNX, and ONNX Runtime.",[50,54,55],{},"Organizing your dataset into train, validation, and test folders for each class.",[50,57,58],{},"Perform \"transfer learning\" to fine-tune the model against your images & generate the ONNX model.",[10,60,61],{},"Let's get started...",[14,63,65],{"id":64},"setup-the-environment","Setup the environment",[67,68,70],"h4",{"id":69},"pre-requisites","Pre-requisites",[10,72,73,74,78],{},"This tutorial was tested on Ubuntu using Python 3 and ",[75,76,77],"code",{},"pyenv"," for environment management.",[10,80,81,82,84],{},"For the sake of brevity, from this point forward, the tutorial will assume you are using a debian based operating system and ",[75,83,77],{},".\nInstructions will need to be adapted if you are using something else.",[86,87,89],"h5",{"id":88},"python-tools","Python tools",[10,91,92,93,95,96,99],{},"Ensure you have ",[75,94,77],{}," and ",[75,97,98],{},"pyenv-virtualenv"," installed.",[101,102,107],"pre",{"className":103,"code":104,"language":105,"meta":106,"style":106},"language-bash shiki shiki-themes github-light github-dark","pyenv --version\npyenv virtualenv --version\n","bash","",[75,108,109,121],{"__ignoreMap":106},[110,111,114,117],"span",{"class":112,"line":113},"line",1,[110,115,77],{"class":116},"sScJk",[110,118,120],{"class":119},"sj4cs"," --version\n",[110,122,124,126,130],{"class":112,"line":123},2,[110,125,77],{"class":116},[110,127,129],{"class":128},"sZZnC"," virtualenv",[110,131,120],{"class":119},[10,133,134,135,140],{},"If you don't have them installed, this ",[25,136,139],{"href":137,"rel":138},"https:\u002F\u002Fmedium.com\u002F@aashari\u002Feasy-to-follow-guide-of-how-to-install-pyenv-on-ubuntu-a3730af8d7f0",[29],"Medium article"," worked well in our case.",[86,142,144],{"id":143},"sub-dependencies","Sub dependencies",[10,146,147,148,151,152,155],{},"During setup and testing, my installation failed at the last step due to missing ",[75,149,150],{},"bz2"," support (a TorchVision dependency).\nIf you encounter this, you would need to install ",[75,153,154],{},"libbz2"," then you would need to rebuild your python environment.\nTo save time, I recommend that you perform the steps below now to ensure the dependencies are installed and avoid the mis-step.",[101,157,159],{"className":103,"code":158,"language":105,"meta":106,"style":106},"sudo apt update\nsudo apt install -y libbz2-dev liblzma-dev libsqlite3-dev libssl-dev zlib1g-dev libffi-dev build-essential\n",[75,160,161,172],{"__ignoreMap":106},[110,162,163,166,169],{"class":112,"line":113},[110,164,165],{"class":116},"sudo",[110,167,168],{"class":128}," apt",[110,170,171],{"class":128}," update\n",[110,173,174,176,178,181,184,187,190,193,196,199,202],{"class":112,"line":123},[110,175,165],{"class":116},[110,177,168],{"class":128},[110,179,180],{"class":128}," install",[110,182,183],{"class":119}," -y",[110,185,186],{"class":128}," libbz2-dev",[110,188,189],{"class":128}," liblzma-dev",[110,191,192],{"class":128}," libsqlite3-dev",[110,194,195],{"class":128}," libssl-dev",[110,197,198],{"class":128}," zlib1g-dev",[110,200,201],{"class":128}," libffi-dev",[110,203,204],{"class":128}," build-essential\n",[67,206,208],{"id":207},"virtual-environment-setup","Virtual Environment Setup",[10,210,211],{},"Install python 3.10.14 (or any version compatible with pytorch and onnx):",[101,213,215],{"className":103,"code":214,"language":105,"meta":106,"style":106},"pyenv install 3.10.14\n",[75,216,217],{"__ignoreMap":106},[110,218,219,221,223],{"class":112,"line":113},[110,220,77],{"class":116},[110,222,180],{"class":128},[110,224,225],{"class":119}," 3.10.14\n",[10,227,228],{},"Create a new virtual environment:",[101,230,232],{"className":103,"code":231,"language":105,"meta":106,"style":106},"pyenv virtualenv 3.10.14 venv_py3_10_14_pytorch\n",[75,233,234],{"__ignoreMap":106},[110,235,236,238,240,243],{"class":112,"line":113},[110,237,77],{"class":116},[110,239,129],{"class":128},[110,241,242],{"class":119}," 3.10.14",[110,244,245],{"class":128}," venv_py3_10_14_pytorch\n",[10,247,248],{},"Activate the virtual environment:",[101,250,252],{"className":103,"code":251,"language":105,"meta":106,"style":106},"pyenv activate venv_py3_10_14_pytorch\n",[75,253,254],{"__ignoreMap":106},[110,255,256,258,261],{"class":112,"line":113},[110,257,77],{"class":116},[110,259,260],{"class":128}," activate",[110,262,245],{"class":128},[10,264,265],{},[266,267,268],"em",{},"NOTE: Depending on your shell, your commandline may become decorated with the name of the virtual environment.",[10,270,271],{},"Install the required packages:",[101,273,275],{"className":103,"code":274,"language":105,"meta":106,"style":106},"pip install --upgrade pip\npip install torch torchvision onnx onnxruntime matplotlib numpy\n",[75,276,277,290],{"__ignoreMap":106},[110,278,279,282,284,287],{"class":112,"line":113},[110,280,281],{"class":116},"pip",[110,283,180],{"class":128},[110,285,286],{"class":119}," --upgrade",[110,288,289],{"class":128}," pip\n",[110,291,292,294,296,299,302,305,308,311],{"class":112,"line":123},[110,293,281],{"class":116},[110,295,180],{"class":128},[110,297,298],{"class":128}," torch",[110,300,301],{"class":128}," torchvision",[110,303,304],{"class":128}," onnx",[110,306,307],{"class":128}," onnxruntime",[110,309,310],{"class":128}," matplotlib",[110,312,313],{"class":128}," numpy\n",[10,315,316],{},"Create a working directory",[101,318,320],{"className":103,"code":319,"language":105,"meta":106,"style":106},"mkdir ~\u002Fmy-py-projects\ncd ~\u002Fmy-py-projects\nmkdir pytorch-onnx\ncd pytorch-onnx\n# Associate this directory with the virtual env we created earlier\npyenv local venv_py3_10_14_pytorch\n",[75,321,322,330,337,345,352,359],{"__ignoreMap":106},[110,323,324,327],{"class":112,"line":113},[110,325,326],{"class":116},"mkdir",[110,328,329],{"class":128}," ~\u002Fmy-py-projects\n",[110,331,332,335],{"class":112,"line":123},[110,333,334],{"class":119},"cd",[110,336,329],{"class":128},[110,338,340,342],{"class":112,"line":339},3,[110,341,326],{"class":116},[110,343,344],{"class":128}," pytorch-onnx\n",[110,346,348,350],{"class":112,"line":347},4,[110,349,334],{"class":119},[110,351,344],{"class":128},[110,353,355],{"class":112,"line":354},5,[110,356,358],{"class":357},"sJ8bj","# Associate this directory with the virtual env we created earlier\n",[110,360,362,364,367],{"class":112,"line":361},6,[110,363,77],{"class":116},[110,365,366],{"class":128}," local",[110,368,245],{"class":128},[10,370,371,372,375],{},"(Optional) Create a ",[75,373,374],{},"requirements.txt"," file to document the packages used in this project:",[101,377,379],{"className":103,"code":378,"language":105,"meta":106,"style":106},"pip freeze > requirements.txt\n",[75,380,381],{"__ignoreMap":106},[110,382,383,385,388,392],{"class":112,"line":113},[110,384,281],{"class":116},[110,386,387],{"class":128}," freeze",[110,389,391],{"class":390},"szBVR"," >",[110,393,394],{"class":128}," requirements.txt\n",[14,396,398],{"id":397},"organizing-your-dataset","Organizing your dataset",[10,400,401,402,407,408,412],{},"For this example, I have created a simple dataset of images of apples, kiwis, and mangos.\nYou can use your own dataset or download a dataset from the internet (e.g. ",[25,403,406],{"href":404,"rel":405},"https:\u002F\u002Fwww.kaggle.com\u002Fdatasets\u002F",[29],"this one"," or ",[25,409,406],{"href":410,"rel":411},"https:\u002F\u002Fimages.cv\u002Fsearch-labeled-image-dataset",[29],").\nJust make sure to organize the images in the following structure:",[101,414,416],{"className":103,"code":415,"language":105,"meta":106,"style":106},"data\u002F\n    train\u002F\n        apples\u002F\n            apple1.jpg\n            apple2.jpg\n            ...\n        kiwis\u002F\n            kiwi1.jpg\n            kiwi2.jpg\n            ...\n        mangos\u002F\n            mango1.jpg\n            mango2.jpg\n            ...\n    val\u002F\n        apples\u002F\n            apple3.jpg\n            apple4.jpg\n            ...\n        kiwis\u002F\n            kiwi3.jpg\n            kiwi4.jpg\n            ...\n        mangos\u002F\n            mango3.jpg\n            mango4.jpg\n            ...\n    test\u002F\n        apples\u002F\n            apple5.jpg\n            apple6.jpg\n            ...\n        kiwis\u002F\n            kiwi5.jpg\n            kiwi6.jpg\n            ...\n        mangos\u002F\n            mango5.jpg\n            mango6.jpg\n            ...\n",[75,417,418,423,428,433,438,443,448,454,460,466,471,477,483,489,494,500,505,511,517,522,527,533,539,544,549,555,561,566,572,577,583,589,594,599,605,611,616,621,627,633],{"__ignoreMap":106},[110,419,420],{"class":112,"line":113},[110,421,422],{"class":116},"data\u002F\n",[110,424,425],{"class":112,"line":123},[110,426,427],{"class":116},"    train\u002F\n",[110,429,430],{"class":112,"line":339},[110,431,432],{"class":116},"        apples\u002F\n",[110,434,435],{"class":112,"line":347},[110,436,437],{"class":116},"            apple1.jpg\n",[110,439,440],{"class":112,"line":354},[110,441,442],{"class":116},"            apple2.jpg\n",[110,444,445],{"class":112,"line":361},[110,446,447],{"class":119},"            ...\n",[110,449,451],{"class":112,"line":450},7,[110,452,453],{"class":116},"        kiwis\u002F\n",[110,455,457],{"class":112,"line":456},8,[110,458,459],{"class":116},"            kiwi1.jpg\n",[110,461,463],{"class":112,"line":462},9,[110,464,465],{"class":116},"            kiwi2.jpg\n",[110,467,469],{"class":112,"line":468},10,[110,470,447],{"class":119},[110,472,474],{"class":112,"line":473},11,[110,475,476],{"class":116},"        mangos\u002F\n",[110,478,480],{"class":112,"line":479},12,[110,481,482],{"class":116},"            mango1.jpg\n",[110,484,486],{"class":112,"line":485},13,[110,487,488],{"class":116},"            mango2.jpg\n",[110,490,492],{"class":112,"line":491},14,[110,493,447],{"class":119},[110,495,497],{"class":112,"line":496},15,[110,498,499],{"class":116},"    val\u002F\n",[110,501,503],{"class":112,"line":502},16,[110,504,432],{"class":116},[110,506,508],{"class":112,"line":507},17,[110,509,510],{"class":116},"            apple3.jpg\n",[110,512,514],{"class":112,"line":513},18,[110,515,516],{"class":116},"            apple4.jpg\n",[110,518,520],{"class":112,"line":519},19,[110,521,447],{"class":119},[110,523,525],{"class":112,"line":524},20,[110,526,453],{"class":116},[110,528,530],{"class":112,"line":529},21,[110,531,532],{"class":116},"            kiwi3.jpg\n",[110,534,536],{"class":112,"line":535},22,[110,537,538],{"class":116},"            kiwi4.jpg\n",[110,540,542],{"class":112,"line":541},23,[110,543,447],{"class":119},[110,545,547],{"class":112,"line":546},24,[110,548,476],{"class":116},[110,550,552],{"class":112,"line":551},25,[110,553,554],{"class":116},"            mango3.jpg\n",[110,556,558],{"class":112,"line":557},26,[110,559,560],{"class":116},"            mango4.jpg\n",[110,562,564],{"class":112,"line":563},27,[110,565,447],{"class":119},[110,567,569],{"class":112,"line":568},28,[110,570,571],{"class":116},"    test\u002F\n",[110,573,575],{"class":112,"line":574},29,[110,576,432],{"class":116},[110,578,580],{"class":112,"line":579},30,[110,581,582],{"class":116},"            apple5.jpg\n",[110,584,586],{"class":112,"line":585},31,[110,587,588],{"class":116},"            apple6.jpg\n",[110,590,592],{"class":112,"line":591},32,[110,593,447],{"class":119},[110,595,597],{"class":112,"line":596},33,[110,598,453],{"class":116},[110,600,602],{"class":112,"line":601},34,[110,603,604],{"class":116},"            kiwi5.jpg\n",[110,606,608],{"class":112,"line":607},35,[110,609,610],{"class":116},"            kiwi6.jpg\n",[110,612,614],{"class":112,"line":613},36,[110,615,447],{"class":119},[110,617,619],{"class":112,"line":618},37,[110,620,476],{"class":116},[110,622,624],{"class":112,"line":623},38,[110,625,626],{"class":116},"            mango5.jpg\n",[110,628,630],{"class":112,"line":629},39,[110,631,632],{"class":116},"            mango6.jpg\n",[110,634,636],{"class":112,"line":635},40,[110,637,447],{"class":119},[10,639,640,641,644],{},"Now, inside ",[75,642,643],{},"~\u002Fmy-py-projects\u002Fpytorch-onnx\u002F"," you should have:",[101,646,648],{"className":103,"code":647,"language":105,"meta":106,"style":106},"pytorch-onnx\u002F\n│\n├── data\u002F\n│   ├── train\u002F\n│   │   ├── apples\u002F\n│   │   ├── kiwis\u002F\n│   │   └── mangos\u002F\n│   ├── val\u002F\n│   └── test\u002F\n│\n├── fruit_classifier.py   # we will create this shortly\n└── requirements.txt      # optional\n",[75,649,650,655,660,668,679,691,702,714,723,732,736,746],{"__ignoreMap":106},[110,651,652],{"class":112,"line":113},[110,653,654],{"class":116},"pytorch-onnx\u002F\n",[110,656,657],{"class":112,"line":123},[110,658,659],{"class":116},"│\n",[110,661,662,665],{"class":112,"line":339},[110,663,664],{"class":116},"├──",[110,666,667],{"class":128}," data\u002F\n",[110,669,670,673,676],{"class":112,"line":347},[110,671,672],{"class":116},"│",[110,674,675],{"class":128},"   ├──",[110,677,678],{"class":128}," train\u002F\n",[110,680,681,683,686,688],{"class":112,"line":354},[110,682,672],{"class":116},[110,684,685],{"class":128},"   │",[110,687,675],{"class":128},[110,689,690],{"class":128}," apples\u002F\n",[110,692,693,695,697,699],{"class":112,"line":361},[110,694,672],{"class":116},[110,696,685],{"class":128},[110,698,675],{"class":128},[110,700,701],{"class":128}," kiwis\u002F\n",[110,703,704,706,708,711],{"class":112,"line":450},[110,705,672],{"class":116},[110,707,685],{"class":128},[110,709,710],{"class":128},"   └──",[110,712,713],{"class":128}," mangos\u002F\n",[110,715,716,718,720],{"class":112,"line":456},[110,717,672],{"class":116},[110,719,675],{"class":128},[110,721,722],{"class":128}," val\u002F\n",[110,724,725,727,729],{"class":112,"line":462},[110,726,672],{"class":116},[110,728,710],{"class":128},[110,730,731],{"class":128}," test\u002F\n",[110,733,734],{"class":112,"line":468},[110,735,659],{"class":116},[110,737,738,740,743],{"class":112,"line":473},[110,739,664],{"class":116},[110,741,742],{"class":128}," fruit_classifier.py",[110,744,745],{"class":357},"   # we will create this shortly\n",[110,747,748,751,754],{"class":112,"line":479},[110,749,750],{"class":116},"└──",[110,752,753],{"class":128}," requirements.txt",[110,755,756],{"class":357},"      # optional\n",[14,758,760],{"id":759},"fine-tune-the-model","Fine-tune the model",[10,762,763],{},"Now we can create a simple pytorch model to classify the images.",[10,765,766,767],{},"Create a new file called ",[75,768,769],{},"fruit_classifier.py",[101,771,773],{"className":103,"code":772,"language":105,"meta":106,"style":106},"# Use nano to create the file (you can use your favorite editor e.g. vim, code, etc)\nnano fruit_classifier.py\n",[75,774,775,780],{"__ignoreMap":106},[110,776,777],{"class":112,"line":113},[110,778,779],{"class":357},"# Use nano to create the file (you can use your favorite editor e.g. vim, code, etc)\n",[110,781,782,785],{"class":112,"line":123},[110,783,784],{"class":116},"nano",[110,786,787],{"class":128}," fruit_classifier.py\n",[10,789,790],{},"Add the following code:",[101,792,796],{"className":793,"code":794,"language":795,"meta":106,"style":106},"language-python shiki shiki-themes github-light github-dark","# fruit_classifier.py\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\nimport onnxruntime as ort\nimport numpy as np\n\n# --- Dataset ---\ndata_dir = \"data\"\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                         std=[0.229, 0.224, 0.225])\n])\n\ntrain_dataset = datasets.ImageFolder(f\"{data_dir}\u002Ftrain\", transform=transform)\nval_dataset   = datasets.ImageFolder(f\"{data_dir}\u002Fval\", transform=transform)\ntest_dataset  = datasets.ImageFolder(f\"{data_dir}\u002Ftest\", transform=transform)\n\ntrain_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\nval_loader   = DataLoader(val_dataset, batch_size=32, shuffle=False)\ntest_loader  = DataLoader(test_dataset, batch_size=32, shuffle=False)\n\nprint(\"Class mapping:\", train_dataset.class_to_idx)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# --- Model ---\nmodel = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)\nmodel.fc = nn.Linear(model.fc.in_features, len(train_dataset.classes))\nmodel = model.to(device)\n\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-4)\n\n\n# --- Training ---\ndef train(num_epochs=5):\n    for epoch in range(num_epochs):\n        model.train()\n        running_loss = 0.0\n        for inputs, labels in train_loader:\n            inputs, labels = inputs.to(device), labels.to(device)\n\n            optimizer.zero_grad()\n            outputs = model(inputs)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n\n            running_loss += loss.item()\n\n        avg_loss = running_loss \u002F len(train_loader)\n        print(f\"Epoch {epoch+1}, Loss: {avg_loss:.4f}\")\n\n\n# --- Evaluation ---\ndef evaluate(loader):\n    model.eval()\n    correct, total = 0, 0\n    with torch.no_grad():\n        for inputs, labels in loader:\n            inputs, labels = inputs.to(device), labels.to(device)\n            outputs = model(inputs)\n            _, preds = torch.max(outputs, 1)\n            correct += (preds == labels).sum().item()\n            total += labels.size(0)\n    return correct \u002F total\n\n\n# --- Export to ONNX ---\ndef export_model():\n    dummy_input = torch.randn(1, 3, 224, 224, device=device)\n\n    torch.onnx.export(\n        model,               # model being run\n        dummy_input,         # model input (or a tuple for multiple inputs)\n        \"fruit_classifier.onnx\",    # where to save the model (can be a file or file-like object)\n        export_params=True,  # store the trained parameter weights inside the model file\n        opset_version=16,    # the ONNX version to export the model to\n        do_constant_folding=True,  # whether to execute constant folding for optimization\n        input_names=['input'],   # the model's input names\n        output_names=['output'],  # the model's output names\n        dynamic_axes={\"input\": {0: \"batch_size\"}, \"output\": {0: \"batch_size\"}}\n    )\n\n    print(\"Model exported to fruit_classifier.onnx\")\n\n\n# --- Test with ONNX Runtime ---\ndef test_onnx():\n    ort_session = ort.InferenceSession(\"fruit_classifier.onnx\")\n\n    def to_numpy(tensor):\n        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n\n    inputs, _ = next(iter(test_loader))\n    ort_inputs = {\"input\": to_numpy(inputs[:1])}\n    ort_outs = ort_session.run(None, ort_inputs)\n\n    pred_class = np.argmax(ort_outs[0])\n    print(\"ONNX Prediction:\", train_dataset.classes[pred_class])\n\n\n# --- Main ---\nif __name__ == \"__main__\":\n    train(num_epochs=5)\n    val_acc = evaluate(val_loader)\n    print(f\"Validation Accuracy: {val_acc:.2%}\")\n\n    export_model()\n    test_onnx()\n\n","python",[75,797,798,803,809,814,819,824,829,834,839,844,849,854,858,863,868,873,878,883,888,893,898,902,907,912,917,921,926,931,936,940,945,949,954,958,963,968,973,978,982,987,992,997,1002,1008,1014,1020,1026,1032,1038,1044,1049,1055,1061,1067,1073,1079,1084,1090,1095,1101,1107,1112,1117,1123,1129,1135,1141,1147,1153,1158,1163,1169,1175,1181,1187,1192,1197,1203,1209,1215,1220,1226,1232,1238,1244,1250,1256,1262,1268,1274,1280,1286,1291,1297,1302,1307,1313,1319,1325,1330,1336,1342,1347,1353,1359,1365,1370,1376,1382,1387,1392,1398,1404,1410,1416,1422,1427,1433],{"__ignoreMap":106},[110,799,800],{"class":112,"line":113},[110,801,802],{},"# fruit_classifier.py\n",[110,804,805],{"class":112,"line":123},[110,806,808],{"emptyLinePlaceholder":807},true,"\n",[110,810,811],{"class":112,"line":339},[110,812,813],{},"import torch\n",[110,815,816],{"class":112,"line":347},[110,817,818],{},"import torch.nn as nn\n",[110,820,821],{"class":112,"line":354},[110,822,823],{},"import torch.optim as optim\n",[110,825,826],{"class":112,"line":361},[110,827,828],{},"from torch.utils.data import DataLoader\n",[110,830,831],{"class":112,"line":450},[110,832,833],{},"import torchvision.transforms as transforms\n",[110,835,836],{"class":112,"line":456},[110,837,838],{},"import torchvision.datasets as datasets\n",[110,840,841],{"class":112,"line":462},[110,842,843],{},"import torchvision.models as models\n",[110,845,846],{"class":112,"line":468},[110,847,848],{},"import onnxruntime as ort\n",[110,850,851],{"class":112,"line":473},[110,852,853],{},"import numpy as np\n",[110,855,856],{"class":112,"line":479},[110,857,808],{"emptyLinePlaceholder":807},[110,859,860],{"class":112,"line":485},[110,861,862],{},"# --- Dataset ---\n",[110,864,865],{"class":112,"line":491},[110,866,867],{},"data_dir = \"data\"\n",[110,869,870],{"class":112,"line":496},[110,871,872],{},"transform = transforms.Compose([\n",[110,874,875],{"class":112,"line":502},[110,876,877],{},"    transforms.Resize((224, 224)),\n",[110,879,880],{"class":112,"line":507},[110,881,882],{},"    transforms.ToTensor(),\n",[110,884,885],{"class":112,"line":513},[110,886,887],{},"    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",[110,889,890],{"class":112,"line":519},[110,891,892],{},"                         std=[0.229, 0.224, 0.225])\n",[110,894,895],{"class":112,"line":524},[110,896,897],{},"])\n",[110,899,900],{"class":112,"line":529},[110,901,808],{"emptyLinePlaceholder":807},[110,903,904],{"class":112,"line":535},[110,905,906],{},"train_dataset = datasets.ImageFolder(f\"{data_dir}\u002Ftrain\", transform=transform)\n",[110,908,909],{"class":112,"line":541},[110,910,911],{},"val_dataset   = datasets.ImageFolder(f\"{data_dir}\u002Fval\", transform=transform)\n",[110,913,914],{"class":112,"line":546},[110,915,916],{},"test_dataset  = datasets.ImageFolder(f\"{data_dir}\u002Ftest\", transform=transform)\n",[110,918,919],{"class":112,"line":551},[110,920,808],{"emptyLinePlaceholder":807},[110,922,923],{"class":112,"line":557},[110,924,925],{},"train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",[110,927,928],{"class":112,"line":563},[110,929,930],{},"val_loader   = DataLoader(val_dataset, batch_size=32, shuffle=False)\n",[110,932,933],{"class":112,"line":568},[110,934,935],{},"test_loader  = DataLoader(test_dataset, batch_size=32, shuffle=False)\n",[110,937,938],{"class":112,"line":574},[110,939,808],{"emptyLinePlaceholder":807},[110,941,942],{"class":112,"line":579},[110,943,944],{},"print(\"Class mapping:\", train_dataset.class_to_idx)\n",[110,946,947],{"class":112,"line":585},[110,948,808],{"emptyLinePlaceholder":807},[110,950,951],{"class":112,"line":591},[110,952,953],{},"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",[110,955,956],{"class":112,"line":596},[110,957,808],{"emptyLinePlaceholder":807},[110,959,960],{"class":112,"line":601},[110,961,962],{},"# --- Model ---\n",[110,964,965],{"class":112,"line":607},[110,966,967],{},"model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)\n",[110,969,970],{"class":112,"line":613},[110,971,972],{},"model.fc = nn.Linear(model.fc.in_features, len(train_dataset.classes))\n",[110,974,975],{"class":112,"line":618},[110,976,977],{},"model = model.to(device)\n",[110,979,980],{"class":112,"line":623},[110,981,808],{"emptyLinePlaceholder":807},[110,983,984],{"class":112,"line":629},[110,985,986],{},"criterion = nn.CrossEntropyLoss()\n",[110,988,989],{"class":112,"line":635},[110,990,991],{},"optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",[110,993,995],{"class":112,"line":994},41,[110,996,808],{"emptyLinePlaceholder":807},[110,998,1000],{"class":112,"line":999},42,[110,1001,808],{"emptyLinePlaceholder":807},[110,1003,1005],{"class":112,"line":1004},43,[110,1006,1007],{},"# --- Training ---\n",[110,1009,1011],{"class":112,"line":1010},44,[110,1012,1013],{},"def train(num_epochs=5):\n",[110,1015,1017],{"class":112,"line":1016},45,[110,1018,1019],{},"    for epoch in range(num_epochs):\n",[110,1021,1023],{"class":112,"line":1022},46,[110,1024,1025],{},"        model.train()\n",[110,1027,1029],{"class":112,"line":1028},47,[110,1030,1031],{},"        running_loss = 0.0\n",[110,1033,1035],{"class":112,"line":1034},48,[110,1036,1037],{},"        for inputs, labels in train_loader:\n",[110,1039,1041],{"class":112,"line":1040},49,[110,1042,1043],{},"            inputs, labels = inputs.to(device), labels.to(device)\n",[110,1045,1047],{"class":112,"line":1046},50,[110,1048,808],{"emptyLinePlaceholder":807},[110,1050,1052],{"class":112,"line":1051},51,[110,1053,1054],{},"            optimizer.zero_grad()\n",[110,1056,1058],{"class":112,"line":1057},52,[110,1059,1060],{},"            outputs = model(inputs)\n",[110,1062,1064],{"class":112,"line":1063},53,[110,1065,1066],{},"            loss = criterion(outputs, labels)\n",[110,1068,1070],{"class":112,"line":1069},54,[110,1071,1072],{},"            loss.backward()\n",[110,1074,1076],{"class":112,"line":1075},55,[110,1077,1078],{},"            optimizer.step()\n",[110,1080,1082],{"class":112,"line":1081},56,[110,1083,808],{"emptyLinePlaceholder":807},[110,1085,1087],{"class":112,"line":1086},57,[110,1088,1089],{},"            running_loss += loss.item()\n",[110,1091,1093],{"class":112,"line":1092},58,[110,1094,808],{"emptyLinePlaceholder":807},[110,1096,1098],{"class":112,"line":1097},59,[110,1099,1100],{},"        avg_loss = running_loss \u002F len(train_loader)\n",[110,1102,1104],{"class":112,"line":1103},60,[110,1105,1106],{},"        print(f\"Epoch {epoch+1}, Loss: {avg_loss:.4f}\")\n",[110,1108,1110],{"class":112,"line":1109},61,[110,1111,808],{"emptyLinePlaceholder":807},[110,1113,1115],{"class":112,"line":1114},62,[110,1116,808],{"emptyLinePlaceholder":807},[110,1118,1120],{"class":112,"line":1119},63,[110,1121,1122],{},"# --- Evaluation ---\n",[110,1124,1126],{"class":112,"line":1125},64,[110,1127,1128],{},"def evaluate(loader):\n",[110,1130,1132],{"class":112,"line":1131},65,[110,1133,1134],{},"    model.eval()\n",[110,1136,1138],{"class":112,"line":1137},66,[110,1139,1140],{},"    correct, total = 0, 0\n",[110,1142,1144],{"class":112,"line":1143},67,[110,1145,1146],{},"    with torch.no_grad():\n",[110,1148,1150],{"class":112,"line":1149},68,[110,1151,1152],{},"        for inputs, labels in loader:\n",[110,1154,1156],{"class":112,"line":1155},69,[110,1157,1043],{},[110,1159,1161],{"class":112,"line":1160},70,[110,1162,1060],{},[110,1164,1166],{"class":112,"line":1165},71,[110,1167,1168],{},"            _, preds = torch.max(outputs, 1)\n",[110,1170,1172],{"class":112,"line":1171},72,[110,1173,1174],{},"            correct += (preds == labels).sum().item()\n",[110,1176,1178],{"class":112,"line":1177},73,[110,1179,1180],{},"            total += labels.size(0)\n",[110,1182,1184],{"class":112,"line":1183},74,[110,1185,1186],{},"    return correct \u002F total\n",[110,1188,1190],{"class":112,"line":1189},75,[110,1191,808],{"emptyLinePlaceholder":807},[110,1193,1195],{"class":112,"line":1194},76,[110,1196,808],{"emptyLinePlaceholder":807},[110,1198,1200],{"class":112,"line":1199},77,[110,1201,1202],{},"# --- Export to ONNX ---\n",[110,1204,1206],{"class":112,"line":1205},78,[110,1207,1208],{},"def export_model():\n",[110,1210,1212],{"class":112,"line":1211},79,[110,1213,1214],{},"    dummy_input = torch.randn(1, 3, 224, 224, device=device)\n",[110,1216,1218],{"class":112,"line":1217},80,[110,1219,808],{"emptyLinePlaceholder":807},[110,1221,1223],{"class":112,"line":1222},81,[110,1224,1225],{},"    torch.onnx.export(\n",[110,1227,1229],{"class":112,"line":1228},82,[110,1230,1231],{},"        model,               # model being run\n",[110,1233,1235],{"class":112,"line":1234},83,[110,1236,1237],{},"        dummy_input,         # model input (or a tuple for multiple inputs)\n",[110,1239,1241],{"class":112,"line":1240},84,[110,1242,1243],{},"        \"fruit_classifier.onnx\",    # where to save the model (can be a file or file-like object)\n",[110,1245,1247],{"class":112,"line":1246},85,[110,1248,1249],{},"        export_params=True,  # store the trained parameter weights inside the model file\n",[110,1251,1253],{"class":112,"line":1252},86,[110,1254,1255],{},"        opset_version=16,    # the ONNX version to export the model to\n",[110,1257,1259],{"class":112,"line":1258},87,[110,1260,1261],{},"        do_constant_folding=True,  # whether to execute constant folding for optimization\n",[110,1263,1265],{"class":112,"line":1264},88,[110,1266,1267],{},"        input_names=['input'],   # the model's input names\n",[110,1269,1271],{"class":112,"line":1270},89,[110,1272,1273],{},"        output_names=['output'],  # the model's output names\n",[110,1275,1277],{"class":112,"line":1276},90,[110,1278,1279],{},"        dynamic_axes={\"input\": {0: \"batch_size\"}, \"output\": {0: \"batch_size\"}}\n",[110,1281,1283],{"class":112,"line":1282},91,[110,1284,1285],{},"    )\n",[110,1287,1289],{"class":112,"line":1288},92,[110,1290,808],{"emptyLinePlaceholder":807},[110,1292,1294],{"class":112,"line":1293},93,[110,1295,1296],{},"    print(\"Model exported to fruit_classifier.onnx\")\n",[110,1298,1300],{"class":112,"line":1299},94,[110,1301,808],{"emptyLinePlaceholder":807},[110,1303,1305],{"class":112,"line":1304},95,[110,1306,808],{"emptyLinePlaceholder":807},[110,1308,1310],{"class":112,"line":1309},96,[110,1311,1312],{},"# --- Test with ONNX Runtime ---\n",[110,1314,1316],{"class":112,"line":1315},97,[110,1317,1318],{},"def test_onnx():\n",[110,1320,1322],{"class":112,"line":1321},98,[110,1323,1324],{},"    ort_session = ort.InferenceSession(\"fruit_classifier.onnx\")\n",[110,1326,1328],{"class":112,"line":1327},99,[110,1329,808],{"emptyLinePlaceholder":807},[110,1331,1333],{"class":112,"line":1332},100,[110,1334,1335],{},"    def to_numpy(tensor):\n",[110,1337,1339],{"class":112,"line":1338},101,[110,1340,1341],{},"        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n",[110,1343,1345],{"class":112,"line":1344},102,[110,1346,808],{"emptyLinePlaceholder":807},[110,1348,1350],{"class":112,"line":1349},103,[110,1351,1352],{},"    inputs, _ = next(iter(test_loader))\n",[110,1354,1356],{"class":112,"line":1355},104,[110,1357,1358],{},"    ort_inputs = {\"input\": to_numpy(inputs[:1])}\n",[110,1360,1362],{"class":112,"line":1361},105,[110,1363,1364],{},"    ort_outs = ort_session.run(None, ort_inputs)\n",[110,1366,1368],{"class":112,"line":1367},106,[110,1369,808],{"emptyLinePlaceholder":807},[110,1371,1373],{"class":112,"line":1372},107,[110,1374,1375],{},"    pred_class = np.argmax(ort_outs[0])\n",[110,1377,1379],{"class":112,"line":1378},108,[110,1380,1381],{},"    print(\"ONNX Prediction:\", train_dataset.classes[pred_class])\n",[110,1383,1385],{"class":112,"line":1384},109,[110,1386,808],{"emptyLinePlaceholder":807},[110,1388,1390],{"class":112,"line":1389},110,[110,1391,808],{"emptyLinePlaceholder":807},[110,1393,1395],{"class":112,"line":1394},111,[110,1396,1397],{},"# --- Main ---\n",[110,1399,1401],{"class":112,"line":1400},112,[110,1402,1403],{},"if __name__ == \"__main__\":\n",[110,1405,1407],{"class":112,"line":1406},113,[110,1408,1409],{},"    train(num_epochs=5)\n",[110,1411,1413],{"class":112,"line":1412},114,[110,1414,1415],{},"    val_acc = evaluate(val_loader)\n",[110,1417,1419],{"class":112,"line":1418},115,[110,1420,1421],{},"    print(f\"Validation Accuracy: {val_acc:.2%}\")\n",[110,1423,1425],{"class":112,"line":1424},116,[110,1426,808],{"emptyLinePlaceholder":807},[110,1428,1430],{"class":112,"line":1429},117,[110,1431,1432],{},"    export_model()\n",[110,1434,1436],{"class":112,"line":1435},118,[110,1437,1438],{},"    test_onnx()\n",[67,1440,1442],{"id":1441},"run-the-fruit_classifierpy-python-script","Run the fruit_classifier.py Python script",[10,1444,1445],{},"Now you can run the script that will train the model, export it to ONNX format, and run a quick classification test using the ONNX Runtime:",[101,1447,1449],{"className":103,"code":1448,"language":105,"meta":106,"style":106},"python fruit_classifier.py\n",[75,1450,1451],{"__ignoreMap":106},[110,1452,1453,1455],{"class":112,"line":113},[110,1454,795],{"class":116},[110,1456,787],{"class":128},[10,1458,1459],{},"What you should see:",[101,1461,1465],{"className":1462,"code":1463,"language":1464,"meta":106,"style":106},"language-log shiki shiki-themes github-light github-dark","Class mapping: {'apple': 0, 'kiwi': 1, 'mango': 2}\nEpoch 1, Loss: 0.7811\nEpoch 2, Loss: 0.1383\nEpoch 3, Loss: 0.0671\nEpoch 4, Loss: 0.0399\nEpoch 5, Loss: 0.0184\nValidation Accuracy: 80.95%\nModel exported to fruit_classifier.onnx\nONNX Prediction: apple\n","log",[75,1466,1467,1472,1477,1482,1487,1492,1497,1502,1507],{"__ignoreMap":106},[110,1468,1469],{"class":112,"line":113},[110,1470,1471],{},"Class mapping: {'apple': 0, 'kiwi': 1, 'mango': 2}\n",[110,1473,1474],{"class":112,"line":123},[110,1475,1476],{},"Epoch 1, Loss: 0.7811\n",[110,1478,1479],{"class":112,"line":339},[110,1480,1481],{},"Epoch 2, Loss: 0.1383\n",[110,1483,1484],{"class":112,"line":347},[110,1485,1486],{},"Epoch 3, Loss: 0.0671\n",[110,1488,1489],{"class":112,"line":354},[110,1490,1491],{},"Epoch 4, Loss: 0.0399\n",[110,1493,1494],{"class":112,"line":361},[110,1495,1496],{},"Epoch 5, Loss: 0.0184\n",[110,1498,1499],{"class":112,"line":450},[110,1500,1501],{},"Validation Accuracy: 80.95%\n",[110,1503,1504],{"class":112,"line":456},[110,1505,1506],{},"Model exported to fruit_classifier.onnx\n",[110,1508,1509],{"class":112,"line":462},[110,1510,1511],{},"ONNX Prediction: apple\n",[10,1513,1514,1515,1518],{},"If you changed the data set from fruit to your use own images and classifications, it will output a different ",[75,1516,1517],{},"Class mapping"," that you will need to use in the Node-RED flow on the next step - make a note of this.",[14,1520,1522],{"id":1521},"using-your-newly-generated-onnx-model-with-the-flowfuse-onnx-node","Using your newly generated ONNX Model with the FlowFuse ONNX Node",[47,1524,1525,1534,1579,1582,1589,1596,1599],{},[50,1526,1527,1528],{},"Import the finished ONNX model into a location in the file system where your Node-RED instance can access it\n",[1529,1530,1531],"ul",{},[50,1532,1533],{},"In FlowFuse cloud you can do this via the Assets tab",[50,1535,1536,1537],{},"Import the demo flow\n",[1529,1538,1539,1542,1553,1561,1568,1574],{},[50,1540,1541],{},"Open your Node-RED editor",[50,1543,1544,1545,1548,1549,1552],{},"Press ",[75,1546,1547],{},"CTRL-I"," or select ",[75,1550,1551],{},"Import"," from the menu to open the Import Dialog",[50,1554,1555,1556,1560],{},"Select the ",[1557,1558,1559],"strong",{},"Examples"," tab",[50,1562,1563,1564,1567],{},"Click the ",[1557,1565,1566],{},"@FlowFuse\u002Fnr-ai-nodes"," entry",[50,1569,1570,1571],{},"Click the demo named ",[1557,1572,1573],{},"advanced-custom-model",[50,1575,1563,1576,1578],{},[1557,1577,1551],{}," Button",[50,1580,1581],{},"Double click the ONNX node to open the configuration dialog",[50,1583,1584,1585,1588],{},"Enter the path to your ONNX model in the ",[1557,1586,1587],{},"Path"," field",[50,1590,1591,1592,1595],{},"If necessary, update the classifications (labels) in the Function node named ",[1557,1593,1594],{},"load labels"," as noted in the previous section",[50,1597,1598],{},"Deploy the flow",[50,1600,1601],{},"Click the inject button on the left of the flow to trigger an inference",[10,1603,1604,1609],{},[1605,1606],"img",{"alt":1607,"dataZoomable":106,"src":1608},"Image showing how to import demo flow","\u002Fblog\u002F2025\u002F10\u002Fimages\u002Fcustom-onnx-mode--import-flow.png",[266,1610,1607],{},[10,1612,1613,1617],{},[1605,1614],{"alt":1615,"dataZoomable":106,"src":1616},"Image showing inference in action","\u002Fblog\u002F2025\u002F10\u002Fimages\u002Fcustom-onnx-mode--in-action.png",[266,1618,1615],{},[14,1620,1622],{"id":1621},"supplementary-notes","Supplementary Notes",[67,1624,1626],{"id":1625},"clean-up","Clean up",[10,1628,1629],{},"To deactivate the virtual environment when you're done, simply run:",[101,1631,1633],{"className":103,"code":1632,"language":105,"meta":106,"style":106},"pyenv deactivate\n",[75,1634,1635],{"__ignoreMap":106},[110,1636,1637,1639],{"class":112,"line":113},[110,1638,77],{"class":116},[110,1640,1641],{"class":128}," deactivate\n",[10,1643,1644,1645,1648],{},"You can remove the ",[75,1646,1647],{},"__pycache__"," and other temporary files if they were created:",[101,1650,1652],{"className":103,"code":1651,"language":105,"meta":106,"style":106},"rm -rf __pycache__\nrm -rf runs\u002F logs\u002F checkpoints\u002F\n",[75,1653,1654,1665],{"__ignoreMap":106},[110,1655,1656,1659,1662],{"class":112,"line":113},[110,1657,1658],{"class":116},"rm",[110,1660,1661],{"class":119}," -rf",[110,1663,1664],{"class":128}," __pycache__\n",[110,1666,1667,1669,1671,1674,1677],{"class":112,"line":123},[110,1668,1658],{"class":116},[110,1670,1661],{"class":119},[110,1672,1673],{"class":128}," runs\u002F",[110,1675,1676],{"class":128}," logs\u002F",[110,1678,1679],{"class":128}," checkpoints\u002F\n",[10,1681,1682],{},"If you want to completely remove the virtual environment, you can do so with:",[101,1684,1686],{"className":103,"code":1685,"language":105,"meta":106,"style":106},"pyenv uninstall venv_py3_10_14_pytorch\n",[75,1687,1688],{"__ignoreMap":106},[110,1689,1690,1692,1695],{"class":112,"line":113},[110,1691,77],{"class":116},[110,1693,1694],{"class":128}," uninstall",[110,1696,245],{"class":128},[1698,1699,1700],"style",{},"html pre.shiki code .sScJk, html code.shiki .sScJk{--shiki-default:#6F42C1;--shiki-dark:#B392F0}html pre.shiki code .sj4cs, html code.shiki .sj4cs{--shiki-default:#005CC5;--shiki-dark:#79B8FF}html pre.shiki code .sZZnC, html code.shiki .sZZnC{--shiki-default:#032F62;--shiki-dark:#9ECBFF}html .default .shiki span {color: var(--shiki-default);background: var(--shiki-default-bg);font-style: var(--shiki-default-font-style);font-weight: var(--shiki-default-font-weight);text-decoration: var(--shiki-default-text-decoration);}html .shiki span {color: var(--shiki-default);background: var(--shiki-default-bg);font-style: var(--shiki-default-font-style);font-weight: var(--shiki-default-font-weight);text-decoration: var(--shiki-default-text-decoration);}html .dark .shiki span {color: var(--shiki-dark);background: var(--shiki-dark-bg);font-style: var(--shiki-dark-font-style);font-weight: var(--shiki-dark-font-weight);text-decoration: var(--shiki-dark-text-decoration);}html.dark .shiki span {color: var(--shiki-dark);background: var(--shiki-dark-bg);font-style: var(--shiki-dark-font-style);font-weight: var(--shiki-dark-font-weight);text-decoration: var(--shiki-dark-text-decoration);}html pre.shiki code .sJ8bj, html code.shiki .sJ8bj{--shiki-default:#6A737D;--shiki-dark:#6A737D}html pre.shiki code .szBVR, html code.shiki .szBVR{--shiki-default:#D73A49;--shiki-dark:#F97583}",{"title":106,"searchDepth":123,"depth":123,"links":1702},[1703,1704,1705,1706,1707,1708,1709,1710],{"id":16,"depth":339,"text":17},{"id":34,"depth":339,"text":35},{"id":41,"depth":339,"text":42},{"id":64,"depth":339,"text":65},{"id":397,"depth":339,"text":398},{"id":759,"depth":339,"text":760},{"id":1521,"depth":339,"text":1522},{"id":1621,"depth":339,"text":1622},"md",{"navTitle":5,"excerpt":1713},{"type":7,"value":1714},[1715],[10,1716,12],{},"\u002Fblog\u002F2025\u002F10\u002Fcustom-onnx-model",{"title":5,"description":12},"blog\u002F2025\u002F10\u002Fcustom-onnx-model","UGbR3G3Xh1_79UofKOOM2uGV0sTci2HbtJR8l7gO-BE",1780070553490]