Experimental results on 2DMambaMIL

[H-Q-N]

Thank you for your excellent work!
About the WSI classification task performed by 2DMambaMIL on the TCGA-BRCA dataset. I ran it three times, and each time the experimental configuration was the default, with epoch=20 and seed=0. The three sets of results I obtained are shown in the figure. As you can see, the results of the three times are different to varying degrees, and you can feel obvious differences in the val dataset and the test dataset. Is this normal? In the results obtained, the value of val_loss will continue to increase with epoch.

[anhtienng]

Hi @H-Q-N ,

Did you use same seed for all runs ?

[H-Q-N]

Hi @H-Q-N ,

Did you use same seed for all runs ?

Yes. The three results are based on the same experimental setup. The results shown are from running the same code three times in a row.

[anhtienng]

Hi @H-Q-N ,
On my side, the curves on validation are normal as follow

I think you could try with the smaller LR first on your side.
I saw your val_loss increases at epoch, I think it's overfitted though I didn't see it on my side.

Regarding the running seed, fixing the seed gives the same results on my side. Maybe your system has uncontrolled randomness.

[H-Q-N]

Hi @H-Q-N , On my side, the curves on validation are normal as follow

I think you could try with the smaller LR first on your side. I saw your val_loss increases at epoch, I think it's overfitted though I didn't see it on my side.

Regarding the running seed, fixing the seed gives the same results on my side. Maybe your system has uncontrolled randomness.

For my experimental results, I modified one function in the mamba_simple.py file, while the other functions remained unchanged.
Specifically, when I ran 2DMambaMIL, an error occurred, located in the reconstruct_2d_wsi function:

def reconstruct_2d_wsi(feats, coords, max_w, max_h, patch_size, pad_token):
    """
    Transform a raster 1D sequence extract by CLAM tool into padded 2D feature
    """
    normalized_coords = coords // patch_size
    w_indexes = normalized_coords[:, 0].int()
    h_indexes = normalized_coords[:, 1].int()
    
    max_w = max_w // patch_size + 1
    max_h = max_h // patch_size + 1
    
    d = feats.shape[-1]
    
    results = pad_token.expand((1, max_h, max_w, d)).to(feats.device).clone()
    results[:, h_indexes, w_indexes, :] = feats

    return results

Traceback (most recent call last):
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/main.py", line 193, in <module>
    main()
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/main.py", line 188, in main
    train(datasets, args)
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/utils/core_utils.py", line 82, in train
    train_loop(epoch, model, train_loader, optimizer, scheduler, loss_fn, args)
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/utils/core_utils.py", line 121, in train_loop
    logits, Y_prob, Y_hat, _, _ = model(data, coords)
                                  ^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/anaconda3/envs/2DMamba/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/anaconda3/envs/2DMamba/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/models/MambaMIL_2D.py", line 86, in forward
    h = self.layers(h, coords, self.pos_embs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/anaconda3/envs/2DMamba/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/anaconda3/envs/2DMamba/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/models/mamba_simple.py", line 101, in forward
    x = layer(x, coords, self.pad_token, pos_embs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/anaconda3/envs/2DMamba/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/anaconda3/envs/2DMamba/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/models/mamba_simple.py", line 123, in forward
    x = reconstruct_2d_wsi(x, coords,
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/gdut-4090/Nian/2DMamba/2DMamba-main/2DMambaMIL/models/mamba_simple.py", line 630, in reconstruct_2d_wsi
    results[:, h_indexes, w_indexes, :] = feats
    ~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Then I added debugging code and got the final function definition：

def reconstruct_2d_wsi(feats, coords, max_w, max_h, patch_size, pad_token):
    """
    Transform a raster 1D sequence extract by CLAM tool into padded 2D feature
    """
    normalized_coords = coords // patch_size
    w_indexes = normalized_coords[:, 0].int()
    h_indexes = normalized_coords[:, 1].int()
    
    max_w = max_w // patch_size + 1
    max_h = max_h // patch_size + 1
    
    d = feats.shape[-1]
    
    results = pad_token.expand((1, max_h, max_w, d)).to(feats.device).clone()


    print("h_indexes:", h_indexes)
    print("w_indexes:", w_indexes)
    print("feats shape:", feats.shape)
    print("results shape:", results.shape)

    # # Check if the index is out of bounds
    # if h_indexes.max() >= results.shape[1]:
    #     print("Error: h_indexes out of bounds. Max h_index:", h_indexes.max().item())
    #     print("Max allowed h_index:", results.shape[1] - 1)
    #     print("Invalid h_indexes:", h_indexes[h_indexes >= results.shape[1]])
    # assert h_indexes.max() < results.shape[1], "h_indexes out of bounds"

    # if w_indexes.max() >= results.shape[2]:
    #     print("Error: w_indexes out of bounds. Max w_index:", w_indexes.max().item())
    #     print("Max allowed w_index:", results.shape[2] - 1)
    #     print("Invalid w_indexes:", w_indexes[w_indexes >= results.shape[2]])
    # assert w_indexes.max() < results.shape[2], "w_indexes out of bounds"

    # Clip h_indexes and w_indexes to ensure their values ​​do not exceed the range
    h_indexes = torch.clamp(h_indexes, 0, results.shape[1] - 1)
    w_indexes = torch.clamp(w_indexes, 0, results.shape[2] - 1)


    # # Check data type
    # assert h_indexes.dtype == torch.int64, "h_indexes must be int64"
    # assert w_indexes.dtype == torch.int64, "w_indexes must be int64"

    # Make sure h_indexes and w_indexes are of type torch.int64
    h_indexes = h_indexes.to(torch.int64)
    w_indexes = w_indexes.to(torch.int64)


    # # Check if the shapes match
    # assert feats.shape == (len(h_indexes), results.shape[-1]), "feats shape mismatch"

    # Reshape feats to remove extra dimensions
    if len(feats.shape) == 3 and feats.shape[0] == 1:
        feats = feats.squeeze(0)  # Reshape from [1, N, C] to [N, C]
    


    results[:, h_indexes, w_indexes, :] = feats

    return results

Is it because I modified this function that the results are so different?

[jingweizhang-xyz]

There are some problems with the indexing issue and we are fixing it.

[ubudde2]

Hi!
I'm also observing some unexplained run-to-run variations when running 2DMambaMIL using --cuda_pscan option. Results are consistent from run to run when not using --cuda_pscan. (Same seed for all runs.) It seems that the backward pass may be the cause of this variability: Logits for the 1st iteration/1st epoch of the forward pass are still identical from run to run, before any grads have been applied to the model parameters. The grads differ, however.

Overall the variations are not particularly large but noticeable. E.g. for 1st epoch on BRACS:

run 1:
TRN 1 Loss: 0.89456 Accuracy: 0.59747 F1: 0.43555 AUC: 0.66604
VAL 1 Loss: 0.69714 Accuracy: 0.73846 F1: 0.54836 AUC: 0.87179
TST 1 Accuracy: 0.65517 F1: 0.50698 AUC: 0.81680

run 2:
TRN 1 Loss: 0.89774 Accuracy: 0.60506 F1: 0.43239 AUC: 0.66035
VAL 1 Loss: 0.69337 Accuracy: 0.73846 F1: 0.54836 AUC: 0.87362
TST 1 Accuracy: 0.65517 F1: 0.50698 AUC: 0.81687

not using --cuda_pscan:
TRN 1 Loss: 0.91008 Accuracy: 0.58734 F1: 0.41913 AUC: 0.64563
VAL 1 Loss: 0.70704 Accuracy: 0.73846 F1: 0.54836 AUC: 0.86375
TST 1 Accuracy: 0.64368 F1: 0.49792 AUC: 0.81521

Is it correct to assume that results with and without --cuda_pscan should be identical?
I'm investigating but am not familiar with cuda code. Any help would be appreciated.
Best regards,
Ulli

[jingweizhang-xyz]

Unlike that on CPU, the logically same programs will not produce the same results on GPUs. As you can see in this VMamba test code https://github.com/MzeroMiko/VMamba/blob/2ed52ead062a51a64521ed3871d52914bf532876/kernels/selective_scan/test_selective_scan.py#L401 The relative difference is set to 6e-4 and absolute difference is set to 2e-3. For gradients on weights, they are all 1e-3. So it is reasonable if the results with and without --cuda_pscan are a bit different. But according to our trials, it does not change our results much.

As to the variability of cuda verion, could you please check the input, output and the returning gradients of this cuda kernel of the two runs both using --cuda_pscan to see if they are the same?

[ubudde2]

I checked input and output for pscan_cuda.fwd (here: https://github.com/AtlasAnalyticsLab/2DMamba/blob/main/2DMambaMIL/models/pscan_2d.py#L40) and they do not change between runs. (1st iteration of 1st epoch only)

Input to pscan_cuda.bwd (here: https://github.com/AtlasAnalyticsLab/2DMamba/blob/main/2DMambaMIL/models/pscan_2d.py#L76) does not change between runs. However output does change. In particular, every element of dA and dB is different.

[jingweizhang-xyz]

Thank you for your testing, we will look into this issue to see if it is a bug.

[jingweizhang-xyz]

Also, could you please send your inputs and dout of this SelectiveScanFn to us (both runs)? So that we can try to reproduce it on our side.

[ubudde2]

Correction: it's actually dB and dC that are different. (indexing error on my part)

[MariosPapadopoulos]

I'm trying to run inference with a WSI embedding of dimension=1024. It's a ResNet50 feature vector from CLAM and I get an OutOfMemory error - how much GPU memory is required to train 2DMambaMIL?

This is my error log:
Traceback (most recent call last):
File "/rds/general/ephemeral/user/kp4718/ephemeral/2DMamba/2DMambaMIL/main.py", line 287, in
main()
File "/rds/general/ephemeral/user/kp4718/ephemeral/2DMamba/2DMambaMIL/main.py", line 264, in main
output = model(input_tensor, coord_tensor)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/rds/general/ephemeral/user/kp4718/ephemeral/2DMamba/2DMambaMIL/models/MambaMIL_2D.py", line 86, in forward
h = self.layers(h, coords, self.pos_embs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/rds/general/ephemeral/user/kp4718/ephemeral/2DMamba/2DMambaMIL/models/mamba_simple.py", line 101, in forward
x = layer(x, coords, self.pad_token, pos_embs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/rds/general/ephemeral/user/kp4718/ephemeral/2DMamba/2DMambaMIL/models/mamba_simple.py", line 142, in forward
output = self.mixer(self.norm(x)) + x
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(args, **kwargs)
File "/rds/general/ephemeral/user/kp4718/ephemeral/2DMamba/2DMambaMIL/models/mamba_simple.py", line 241, in forward
x = self.conv1d(x)[:, :, :WH] # depthwise convolution over time, with a short filter
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 310, in forward
return self._conv_forward(input, self.weight, self.bias)
File "/rds/general/user/kp4718/home/anaconda3/envs/2dmamba/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 306, in _conv_forward
return F.conv1d(input, weight, bias, self.stride,
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 9.32 GiB. GPU 0 has a total capacity of 44.42 GiB of which 6.62 GiB is free. Including non-PyTorch memory, this process has 37.79 GiB memory in use. Of the allocated memory 37.29 GiB is allocated by PyTorch, and 9.25 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)