export_sam_model on Windows not working; RuntimeError: "upsample_bilinear2d_out_frame" not implemented for 'Byte'

[maartenpaul]

Describe the bug
When attempting to export a fine-tuned micro-sam model to the BioImage.IO format using the export_sam_model function on Windows 10, a RuntimeError is raised. This appears to be due to a missing pre-compiled kernel in the PyTorch Windows build for a specific upsampling operation on 'Byte' (uint8) data types. The process works as expected on Linux.

To Reproduce

1. Set up a micro-sam environment on Windows 10 using conda from conda-forge.
2. Install micro-sam version 1.6.1 and a recent PyTorch build (e.g., 2.7.1 with CUDA).
3. Execute a script that calls the export_sam_model function with a uint8 image as input.

Error Log

The script fails with the following RuntimeError:

...
File ...\site-packages\torch\nn\functional.py:4678, in interpolate(input, size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
    4676 assert align_corners is not None
    4677 if antialias:
-> 4678     return torch._C._nn._upsample_bilinear2d_aa(
    4679         input, output_size, align_corners, scale_factors
    4680     )
...
RuntimeError: "upsample_bilinear2d_out_frame" not implemented for 'Byte'

Workaround

I have confirmed that the same script runs without any issues in a Linux environment (WSL Ubuntu) with a comparable PyTorch version. This suggests the necessary kernel for this operation is available in the Linux PyTorch builds but not in the Windows ones. I thought it will be useful report in case somebody hits into the same issue.

Environment

• OS: Windows 10
• pytorch version: 2.7.1 (installed via conda from conda-forge)
• micro-sam version: 1.6.1
• GPU: NVIDIA GeForce RTX 5090
• CUDA version: 12.9

[anwai98]

Hi @maartenpaul,

Thanks for reporting and checking this issue out in detail! To spell out the core suspicion at the moment for other developers:

I have confirmed that the same script runs without any issues in a Linux environment (WSL Ubuntu) with a comparable PyTorch version. This suggests the necessary kernel for this operation is available in the Linux PyTorch builds but not in the Windows ones.

Could we try it with a lower pytorch version? (i.e. install micro-sam in a new environment with conda create -n test-neu -c conda-forge python=3.11 micro_sam "pytorch<2.7"

I'll meanwhile try to search around if I find a hint at what's going on from your error trace!

[anwai98]

Also on a second thought @maartenpaul, could you please the script with me so that I can try to reproduce and see if I can find a workaround for this?

[maartenpaul]

Hi @anwai98 ,
Thanks for the quick reply. I'm running the code in a jupyter notebook. I have tried to lower pytorch<2.7 but I'm getting the same error. If you want a full error log or the finetuned model. It seems to fail on testing the test image.

import numpy as np
from micro_sam.bioimageio.model_export import export_sam_model
import os
from tifffile import imread
import micro_sam.util as util

# Get a test image and label to use for exporting
# For this example, we'll use the first image and label from validation set
test_image_path = os.path.join("D:/Maarten/micro-sam-20250829_134050/val_input", "input_00001.tif")
test_label_path = os.path.join("D:/Maarten/micro-sam-20250829_134050/val_label", "label_00001.tif")
output_directory = "D:/Maarten/micro-sam-20250829_134050"
# Load the test image and label
test_image = imread(test_image_path)
test_label = imread(test_label_path)
model_type="vit_b"
checkpoint_name = "micro_sam_training_20250829_135659"
# Define the path for saving the bioimage.io model
bioimageio_model_path = os.path.join(output_directory, "micro_sam_training_test.zip")

# Export the SAM model to bioimage.io format
export_sam_model(
    image=test_image,
    label_image=test_label,
    model_type=model_type,  # Using the same model_type as in training
    name=f"micro_sam_test",
    output_path=bioimageio_model_path,
    checkpoint_path=os.path.join(
        output_directory, "checkpoints", checkpoint_name, "best.pt"
    ),
    # Optional: Add additional kwargs as needed
    authors=[{"name": "Your Name", "affiliation": "Your Institution"}],
    description="Micro-SAM model trained on microscopy images for segmentation",
    license="MIT",
    documentation="Model trained with micro-sam for segmenting microscopy images",
)

print(f"BioImage.IO model exported to: {bioimageio_model_path}")

[anwai98]

Hi @maartenpaul,

Thanks for sending in the code.

I saw some posts in the PyTorch forum stating (in a bit different context though), this can be fixed by passing the image as floating tensors. Could you try to normalize the image (to convert it away from uint8) and see if that fixes things? (the simplest heuristics would be to divide the image by 255)

Let me know if that works for you!

[maartenpaul]

Hi @anwai98 !
Thanks for the suggestion. I have tried your suggestion, but it seems somewhere in the code the image in converted back to unit8 .

micro-sam/micro_sam/util.py

Line 605 in 589e521

def _to_image(input_):

[anwai98]

Hi @maartenpaul,

Thanks for checking this out, and apologies for the late response!

Can you debug this and see if hard-coding the particular part out helps with your issue? (I'd like to identify the issue - and I am quite sure now that it's the non-floating tensors as I faced it today myself)

If it does, we can find a workable solution for you!

[maartenpaul]

I'll give it a shot in the coming days. I've tried to patch the function but then somewhere further down in the bioimage.io function I ran into issues (that's why likely a 'uint8' was required in the first place).
For now it works for me to export the models in Linux/WSL, but indeee it would be nice to do so in Windows as well. I will try to see if I can help to find a fix.

[maartenpaul]

I managed to export a model on Windows commenting out a few lines of code

First converting the input image to a float

print(type(test_image))
#test_image = test_image/test_image.max()  # normalizing even not neccesary
test_image = test_image.astype("float32")

Commenting out this code in util.py to avoid conversion back to uint8

def _to_image(input_):
    # we require the input to be uint8
   # if input_.dtype != np.dtype("uint8"):
        # first normalize the input to [0, 1]
        input_ = input_.astype("float32") - input_.min()
        input_ = input_ / input_.max()
        # then bring to [0, 255] and cast to uint8
        input_ = (input_ * 255).astype("uint8")

Then got an error "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
". After making some changes (thanks to GTP5) it works for now. See here my changes:
maartenpaul@3458c70
Not a proper fix, but indeed the issue is with the type of the input image. I hope this can help to resolve this issue.

[anwai98]

Hi @maartenpaul,

Thanks for checking it out. And it's good to confirm that it's an image issue! (I was kind of sure about it, but great that you figured out the location where the issue is caused as well).

I'll check this out at some point and try to find a working patch for ensuring floating point inputs!