Voxel-wise classification training pipeline

[relyativist]

I want to train simple neural network two layer classifier. Input data is 3-dimensional volumes of feature maps extracted from Brain MRI. I have manually extracted 6 feature maps each size of [128,128,128]. I want to feed each voxel into neural network, so input layer will have 6 neurons (for each map). The labels are the masks that manually segmented in region of brain malformation. I got following question regarding implementation:

The issue is I don’t quite understand how to create a training pipeline and dataloader for this type of input. I assume to create PatchDataset with size 1x1x1 https://docs.monai.io/en/stable/data.html?highlight=patchdataset#patchdataset. Are there examples of voxel-wise training example?

I apologies, if this question doesn't relate discussion topics.
Thanks

[ebrahimebrahim]

it sounds like there's nothing spatial about the model you want to learn? i.e. the prediction for a given voxel should be independent of the features in the other voxels near to it?

if so then maybe PatchDataset is overkill and all you need to do is reshape to absorb the spatial dimension into the batch dimension

• e.g. if you have a batch size of $B$ and each image consists of 6 feature maps in a $H\times W \times D$ image, then the shape might be $(B,6,H,W,D)$.
• img.permute(0,2,3,4,1) will move the spatial dimensions to be adjacent to batch dimension giving you shape $(B,H,W,D,6)$
• img.reshape(-1,6) will then absorb the spatial dimensions into the batch dimension, giving shape $(BHWD, 6)$, so now it just looks like - the batch size is the product $BHWD$ and each individual data item is a 6-element vector.
• input this into your model during training/inference and then undo the steps above to pull the spatial dimensions back out of the batch dimension to get your output in the form of a 3d image at the end

[ebrahimebrahim]

glad it worked out!

if you want to share then I'm curious to hear how that went for brain lesion segmentation

[aalhayali]

Hey @ebrahimebrahim !

I wanted to hop on to this discussion to ask a question relating to shape adjustments. I posted my issue here (#5372), and I wanted to ask if the input for the training model should be (B,H,W,D,Channel) or (B,Channel,D,H,W). I thought since it follows Pytroch convention, the latter would be the input to the network. The current input I feed into my network is the latter. Would you recommend permuting my image to (B,H,W,D,Channel) ?

Thanks!

[ebrahimebrahim]

@aalhayali hi!

right (B,Channel,D,H,W) is the way to go if you are using pytorch/monai tools.
or (B,Channel,H,W,D) -- I am not differentiating among spatial dimensions H,W,D

in my discussion above when I wrote $(BHWD,6)$ there are two dimensions and the first one has size the product $BHWD$.
so it's still ordered (Batch, Channel), it's just that the batch size is a product

[aalhayali]

@ebrahimebrahim got it!

Quick follow up - when training a 3D U-Net model (in_channel = 1, out_channel = 2) for semantic segmentation, should the input be (B,1,D,H,W) or (B,2,D,H,W)? When using the settings I have for my problem, the shape of the input is (B,1,D,H,W), but since the out_channel = 2, loss calculation generates an error since the original label has a channel size of 1 while prediction is 2. Any suggestions?

[ebrahimebrahim]

Quick follow up - when training a 3D U-Net model (in_channel = 1, out_channel = 2) for semantic segmentation, should the input be (B,1,D,H,W) or (B,2,D,H,W)? When using the settings I have for my problem, the shape of the input is (B,1,D,H,W), but since the out_channel = 2, loss calculation generates an error since the original label has a channel size of 1 while prediction is 2. Any suggestions?

you can post this as a separate discussion topic to get more answers, but from what you've said it sounds like maybe you need to pay attention to whether labels are one-hot encoded / probabilities or whether they are single channel class labels. if you don't want to bother converting things, MONAI loss functions sometimes have some parameter like to_onehot built in to include that step in the loss function call

https://en.wikipedia.org/wiki/One-hot#Machine_learning_and_statistics