Finetuning microsam AIS (continued)

[chuyaowang]

Hello, continuing from this issue #1211, our data also contain some mammalian cells in addition to the hyphal cells. We grow them in a coculture. The mammalian cells are imaged with a fluorescence cytoplasmic channel and a DAPI nuclear channel, while the hyphal cells are best visualized with DIC. When in focus, the mammalian cells can be segmented quite well by AIS, but is a little problematic when a part of a cell goes out of focus while a part of it is still in focus along the z stack. So I also want to fine-tune the model on the mammalian cell z-stack images.

Since the mammalian cell images contain 2 channels. How should I process the channels for finetuning, as the microsam.train._check_loader only accepts 1 or 3 channel inputs? The microsam.util._to_image function just adds a third channel of 0s, but in your paper, you described: "averaging the two channels to obtain a single channel that is then duplicated three times works better ... but is detrimental compared to using both channels independently." I want to align with your training protocol as close as possible to minimize the amount of new data distribution the model needs to learn.

[anwai98]

Hi @chuyaowang,

Thanks for reaching out with a follow-up.

I'll drop a high-level comment before digging into some detailed feedback: please help us out with sharing some images / screenshots so that we can have a closer look at your images, and understand the task better. Otherwise, we are making open guesses and the iteration to reach a solution increases by a lot. I'll give my best shot to answer the questions anyways!

When in focus, the mammalian cells can be segmented quite well by AIS, but is a little problematic when a part of a cell goes out of focus while a part of it is still in focus along the z stack. So I also want to fine-tune the model on the mammalian cell z-stack images.

I see, gotcha!

Since the mammalian cell images contain 2 channels. How should I process the channels for finetuning, as the microsam.train._check_loader only accepts 1 or 3 channel inputs?

Good question. So for this, what we do while training is consider the cytoplasm as channel 1, DAPI as channel 2, and make channel 3 as zeros (eg. the TissueNet data is a good example for such a situation). And if necessary, the first two channels are the only ones which should be normalized to range [0, 255], otherwise leave it as is if it's in the expected range)

I want to align with your training protocol as close as possible to minimize the amount of new data distribution the model needs to learn.

See above, that should do the trick! ;)

Let us know if there are further questions. Hope this helps!

PS. I'll leave this issue open in case more things come up!

[chuyaowang]

Hi, I do have a new issue. During training the AIS decoder, the model does not predict anything in the forward pass. As a result, the losses are NaNs.

I ran a forward pass manually and found that under FP16, the model produces NaNs but under FP32 it is alright.

from micro_sam.training.util import get_trainable_sam_model
from micro_sam.instance_segmentation import get_unetr

device = get_device()
sam_model, state = get_trainable_sam_model(
    model_type="vit_b_lm", device=device,
    checkpoint_path=encoder_path, return_state=True, decoder_path=decoder_path,
)
model = get_unetr(
    image_encoder=sam_model.sam.image_encoder,
    decoder_state=state.get("decoder_state", None),
    device=device,
).eval()

for raw, label in train_loader:
    raw = raw.to(device)
    # test fp32
    with torch.no_grad():
        pred_fp32 = model(raw)
    print("fp32:", pred_fp32.min().item(), pred_fp32.max().item(), "nan:",
torch.isnan(pred_fp32).any().item())
    # test fp16
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
        pred_fp16 = model(raw)
    print("fp16:", pred_fp16.min().item(), pred_fp16.max().item(), "nan:",
torch.isnan(pred_fp16).any().item())
    break

fp32: 7.736363841104321e-06 0.9999996423721313 nan: False
fp16: nan nan nan: True

I noticed during training mixed precision is used automatically, but for me this isn't helping. Is there another way to mitigate the issue? If there isn't I can use all FP32 for training, but this exceeds my memory limit, and I will need to use something like kaggle notebooks for better GPUs.

[anwai98]

Hmm I am a bit lost in text here. Let's step back a bit:

1. Which hardware are you using to train micro-sam?
2. Which pytorch version do you have?
3. Can you try the latest micro-sam version with the example script for instance segmentation? https://github.com/computational-cell-analytics/micro-sam/blob/main/examples/finetuning/finetune_hela_instance_segmentation.py

[chuyaowang]

Hardware:

Pytorch version 2.6.0; microsam version 1.7.7

This is the function call I'm using for training:

from micro_sam.util import get_device
from micro_sam.training import train_instance_segmentation
from pathlib import Path

device = get_device()

save_root = Path("~/Projects/C_Albicans Project/7. Data/model_checkpoints/micro_sam_finetune").expanduser()

encoder_path = Path("~/Projects/C_Albicans Project/7. Data/model_checkpoints/micro_sam/models/vit_b_lm").expanduser()
decoder_path = Path("~/Projects/C_Albicans Project/7. Data/model_checkpoints/micro_sam/models/vit_b_lm_decoder").expanduser()

train_instance_segmentation(
    name="microsam_candida",
    model_type="vit_b_lm",
    train_loader=train_loader,
    val_loader=val_loader,
    checkpoint_path=encoder_path,
    n_epochs=100,
    save_root=save_root,
    device=device,
    freeze=["image_encoder"],
    decoder_path=decoder_path,
)

I also uploaded an image and its label in case you need it.

data.zip

[anwai98]

HI @chuyaowang,

Re: shared image and labels: That's great, thanks. This makes it easy for me to debug and tell you if the training works with the latest setup!

And re: hardware: Hmm as far as my search goes, turing architecture GPUs should support mixed precision. That's confusing.

Let's do it this way: I'll give finetuning a quick shot on your data. It'll be best to try it out after that on some open-source resource to validate if it's your hardware that's fussy or something else.

I'll come back to you in a couple of days!

[chuyaowang]

An update: I have found the cause.

I tried to run the code on kaggle. With the T4 GPU they provide I can finetune the vit-b-lm encoder and decoder at the same time. The training is proceeding normally and I'm seeing the loss going down now, and I am not forcing all FP32 or FP16 at all, just using the vanilla settings. So it does seem like a hardware issue with my 1660 Ti.

Interestingly enough I encountered a similar issue with another deep learning model facefusion (for face swapping) a couple years ago. That model was also giving me all black predictions on my 1660Ti, while normally an image with a swapped face should be predicted.

We had a discord thread about it.

That issue was fixed by using EXHAUSTIVE rather than DEFAULT for cudnn_conv_algo_search in their code.

def apply_execution_provider_options(execution_providers: List[str]) -> List[Any]:
    execution_providers_with_options : List[Any] = []

    for execution_provider in execution_providers:
        if execution_provider == 'CUDAExecutionProvider':
            execution_providers_with_options.append((execution_provider,
            {
                'cudnn_conv_algo_search': 'EXHAUSTIVE'
            }))
        else:
            execution_providers_with_options.append(execution_provider)
    return execution_providers_with_options

The reason for that issue was that some CUDNN convolution algorithms were numerically unstable for FP16. The intermediate outputs would exceed FP16 limits and cause Infs. The Infs would then cause NaNs. So changing it to EXHAUSTIVE search made CUDNN find the fastest algorithm that also happened to be numerically stable to use instead of the numerically unstable default one.

I suspected this could be a similar issue. I looked at where are the NaNs coming from, and the results are

fp32 full  : 1.6765909094829112e-05 0.9999949932098389 nan: False
fp16 enc[0]: max=1.0  nan=False
fp16 full  : nan nan nan: True

So FP32 on the encoder+decoder or FP16 on encoder only are both ok, but FP16 on the decoder caused the NaNs. I further looked at where the NaNs originate in the decoder:

First NaN at: deconv1.block.1.block

  ok   max=       0.9  (1, 256, 64, 64)                     encoder.neck
  ok   max=       0.9  (1, 256, 64, 64)                     encoder
  ok   max=       0.3  (1, 512, 128, 128)                   deconv1.block.0.conv
  ok   max=       0.3  (1, 512, 128, 128)                   deconv1.block.0
  NaN  max=       nan  (1, 512, 128, 128)                   deconv1.block.1.block <<< FIRST NaN
  NaN  max=       nan  (1, 512, 128, 128)                   deconv1.block.1
  NaN  max=       nan  (1, 512, 128, 128)                   deconv1.block.2
  NaN  max=       nan  (1, 512, 128, 128)                   deconv1.block.3
  NaN  max=       nan  (1, 512, 128, 128)                   deconv1.block
  NaN  max=       nan  (1, 512, 128, 128)                   deconv1

This is a 512 x 512 convolution block:

<class 'torch.nn.modules.conv.Conv2d'>
Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
input:  nan=False  max=0.2705
output: nan=True  max=nan

This may point to the same underlying cause. I tested this by setting torch.backends.cudnn.benchmark = True and torch.backends.cudnn.benchmark = False or True. These two settings vary the CUDNN algorithm used by PyTorch, but I still have NaNs at FP16 regardless of which one I use. So this is a hardware limitation. The reason is that 1660Ti does not have tensor cores, only cuda cores. So it can only accumulate intermediate convolutional outputs at FP16, while gpus like T4 have tensor cores and can accumulate them in FP32 while keeping inputs and outputs at FP16. Why facefusion worked by changing the algorithm could be they used smaller convolutions. I'm only guessing as I no longer have the model.

The solution, as I have found, is to use BF16 for computing. It has a higher limit than FP16 and less memory intensive then FP32 at the cost of small precision loss.

fp16  deconv1 Conv2d(512,512,3�3):  in_max=0.3186  out_nan=True  out_max=nan
      full model output:  nan=True
bf16  deconv1 Conv2d(512,512,3�3):  in_max=0.3184  out_nan=False  out_max=0.6719
      full model output:  nan=False

I have opened two PRs for torch-em and micro-sam to implement the fix. The torch-em fix includes a dtype argument to control what dtype to use during training. The micro-sam fix detects the device type and calls the torch-em function with the correct arguments. I have tested them on my GPU and it fixed the NaN issue, but you should test it on other GPUs to make sure nothing is broken.