Cross entropy loss for 3D 1-class segmentation (spleen tutorial)

[ctpn]

Hello, thank you for developing MONAI!

Has anyone gotten the 3D spleen segmentation tutorial working, but using Pytorch's binary cross entropy or cross entropy loss, instead of Dice loss? I have not been successful and would be grateful for any feedback.

To summarize what I think is the problem: Pytorch's BCE loss only works with 1 channel inputs (understandably). However, the 3D spleen segmentation has out_channels=2 in the U-Net, despite being a 1-class segmentation (is there a reason for this?). If I change out_channels=1, then the DiceMetric either stays at 0 or some other constant value throughout training. I thought the issue might have to do with the post-transforms,
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=True, num_classes=2)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=True, num_classes=2)]) ,
...
val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
val_labels = [post_label(i) for i in decollate_batch(val_labels)]
and tried various combinations and permutations. Weirdly, changing post_pred, which should only affect the outputs, changes the shape of val_outputs, even if post_label is kept constant.

E.g.
post_pred = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(argmax=True)])
post_label = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(to_onehot=True, num_classes=2)]) ,
...
val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
val_labels = [post_label(i) for i in decollate_batch(val_labels)]
yields
val_outputs with shape (1, 512, 512, 22)
val_labels with shape(2, 512, 512, 22)

BUT
post_pred = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(argmax=True,**to_onehot=True, num_classes=2**)])
post_label = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(to_onehot=True, num_classes=2)]) ,
...
val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
val_labels = [post_label(i) for i in decollate_batch(val_labels)]
yields
val_outputs with shape (2, 512, 512, 22)
val_labels with shape(1, 512, 512, 22)

i.e. despite post_label being the same in both cases, the shape of val_labels seems to change...?

The rest of the code is basically verbatim the 3D spleen segmentation example. My data is 512 x 512 x [10~30]. I have adjusted pixdim accordingly.

Thanks so much for your time. Any advice would be appreciated.

[Nic-Ma]

Hi @ctpn ,

Thanks for your interest and interesting experiments with MONAI.
About your issue, I feel maybe it's caused by some notebook cache? Could you please help totally restart and rerun the notebook to verify? If still there, please share the test program and error log details for further check.

Thanks.

[ctpn]

Hi @Nic-Ma,

Thanks so much for your quick reply!

What you suggested makes sense to me. After restarting and rerunning the notebook, I've encountered the following RuntimeError (RuntimeError: Sizes of tensors must match except in dimension 4. Got 4 and 3 (The offending index is 0)) instead. I've consistently run into this problem in the past as well, even on code that is almost the same as that in the spleen segmentation tutorial.

code:

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Spacingd(keys=["image", "label"], pixdim=(
            1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        EnsureTyped(keys=["image", "label"]),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Spacingd(keys=["image", "label"], pixdim=(
            1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        EnsureTyped(keys=["image", "label"]),
    ]
)

train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4)
train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=4)  # batch_size=1
val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device: ', device)

model = UNet(
    dimensions=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean")  # why include background? https://github.com/Project-MONAI/MONAI/discussions/1727
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)

max_epochs = 600
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=True, n_classes=2)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=True, n_classes=2)])

for epoch in range(max_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        print(
            f"{step}/{len(train_ds) // train_loader.batch_size}, "
            f"train_loss: {loss.item():.4f}")
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                roi_size = (160, 160, 160)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(
                    val_inputs, roi_size, sw_batch_size, model)
                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
                val_labels = [post_label(i) for i in decollate_batch(val_labels)]
                # compute metric for current iteration
                dice_metric(y_pred=val_outputs, y=val_labels)

            # aggregate the final mean dice result
            metric = dice_metric.aggregate().item()
            # reset the status for next validation round
            dice_metric.reset()

            metric_values.append(metric)
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), os.path.join(
                    root_dir, "best_metric_model.pth"))
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f"\nbest mean dice: {best_metric:.4f} "
                f"at epoch: {best_metric_epoch}"
            )

outputs: (I added print lines in monai.networks.layers.simplelayers.py, class SkipConnection, forward function, ~ line 127. Displayed as 'x: torch.Size... y: torch.Size...')

----------
epoch 1/600
x: torch.Size([1, 128, 32, 32, 2])
y: torch.Size([1, 256, 32, 32, 2])
x: torch.Size([1, 64, 64, 64, 3])
y: torch.Size([1, 64, 64, 64, 4])

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-19-a10e6cb32f70> in <module>
     21         )
     22         optimizer.zero_grad()
---> 23         outputs = model(inputs)
     24         loss = loss_function(outputs, labels)
     25         loss.backward()

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/nets/unet.py in forward(self, x)
    187 
    188     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 189         x = self.model(x)
    190         return x
    191 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    123 
    124     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 125         y = self.submodule(x)
    126 
    127         if self.mode == "cat":

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    123 
    124     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 125         y = self.submodule(x)
    126 
    127         if self.mode == "cat":

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    128             print('x:',x.shape)
    129             print('y:',y.shape)
--> 130             return torch.cat([x, y], dim=self.dim)
    131         if self.mode == "add":
    132             return torch.add(x, y)

RuntimeError: Sizes of tensors must match except in dimension 4. Got 4 and 3 (The offending index is 0)

Please let me know if you have any suggestions, thanks so much!

[Nic-Ma]

Hi @ctpn ,

I feel that because some data in your dataset has different dims, could you please help double confirm all the images in your dataset?
By the way, may I know why you use CTPN network as user name?

Thanks.

[ctpn]

Hi @Nic-Ma, thanks for your suggestion and for the paper! The transform EnsureChannelFirstd is included in train_transforms, would it be possible that it's still not channel-first despite this?

Hi @rijobro, thanks for your suggestion! I tried without the CropForegroundd transform, so that all that was left was:

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Spacingd(keys=["image", "label"], pixdim=(
            1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
        EnsureTyped(keys=["image", "label"]),
    ]
)

And still got the same error:

----------
epoch 1/600
x: torch.Size([1, 128, 22, 22, 1])
y: torch.Size([1, 256, 22, 22, 1])
x: torch.Size([1, 64, 43, 43, 2])
y: torch.Size([1, 64, 44, 44, 2])

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-33-a10e6cb32f70> in <module>
     21         )
     22         optimizer.zero_grad()
---> 23         outputs = model(inputs)
     24         loss = loss_function(outputs, labels)
     25         loss.backward()

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/nets/unet.py in forward(self, x)
    187 
    188     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 189         x = self.model(x)
    190         return x
    191 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    123 
    124     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 125         y = self.submodule(x)
    126 
    127         if self.mode == "cat":

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    123 
    124     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 125         y = self.submodule(x)
    126 
    127         if self.mode == "cat":

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    128             print('x:',x.shape)
    129             print('y:',y.shape)
--> 130             return torch.cat([x, y], dim=self.dim)
    131         if self.mode == "add":
    132             return torch.add(x, y)

RuntimeError: Sizes of tensors must match except in dimension 2. Got 44 and 43 (The offending index is 0)

I then uninstalled and reinstalled MONAI, in case I had accidentally changed some code beyond the print lines I had previously mentioned, and got the same error. Then tried to use a different file, same error. Not sure what I broke... this was fine a few days ago...

----------
epoch 1/600

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-36-a10e6cb32f70> in <module>
     21         )
     22         optimizer.zero_grad()
---> 23         outputs = model(inputs)
     24         loss = loss_function(outputs, labels)
     25         loss.backward()

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/nets/unet.py in forward(self, x)
    187 
    188     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 189         x = self.model(x)
    190         return x
    191 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    123 
    124     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 125         y = self.submodule(x)
    126 
    127         if self.mode == "cat":

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    123 
    124     def forward(self, x: torch.Tensor) -> torch.Tensor:
--> 125         y = self.submodule(x)
    126 
    127         if self.mode == "cat":

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/miniconda3/envs/jupyterhub/lib/python3.9/site-packages/monai/networks/layers/simplelayers.py in forward(self, x)
    126 
    127         if self.mode == "cat":
--> 128             return torch.cat([x, y], dim=self.dim)
    129         if self.mode == "add":
    130             return torch.add(x, y)

RuntimeError: Sizes of tensors must match except in dimension 4. Got 10 and 9 (The offending index is 0)

[rijobro]

If your images are initially all different sizes to each other, then removing the CropForeground will still give you images of different sizes. Could you do something like:

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Spacingd(keys=["image", "label"], pixdim=(
            1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
    ]
)

for i in all_data:
    d = train_transforms[i]
    print(d["image"].shape, d["label"].shape)

and verify that all images are the same size.

[ctpn]

Thanks @rijobro for the suggestion.

I used the train_transform that you suggested verbatim, and modified the for loop a bit as I could not get it to run as-is (sorry):

for i in range(0, len(train_files)):
    d = train_transforms(train_files[i])
    print(d["image"].shape, d["label"].shape)

output:

(1, 512, 512, 22) (1, 512, 512, 22)
(1, 512, 512, 51) (1, 512, 512, 51)
(1, 512, 512, 71) (1, 512, 512, 71)
(1, 512, 512, 51) (1, 512, 512, 51)

Is the problem that the number of Z-slices are different?

Thanks again for your help!

(Unrelated but thank you as well for the bootcamp, it was great to hear MONAI's presentations these past two days!)

[rijobro]

Yes, this will cause problems. Knowing your minimum and maximum sizes in the z direction, you could crop or pad so that they're all the same size.

[ctpn]

Thank you! Not sure if this is relevant to the error but the NiFTi header was also incorrect.