Using this UNet implementation: https://github.com/kilgore92/PyTorch-UNet
With this instantiation of the UNet object:
model = UNet(n_channels=1,
mode='3D',
num_classes=1,
use_pooling=True,
)
Whenever I try running my training script with the dataLoader I get this error: RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 1, 3, 3], but got 5-dimensional input of size [1, 1, 128, 128, 128] instead
Here’s more information:
This is my training script, and it is erroring out at the output = model.forward(images)
.
def train():
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
num_steps_train = len(train_loader)
print(num_steps_train)
for epoch in range(epochs):
print(' - training - ')
for i, (images, masks) in enumerate(train_loader):
images = images.to(device)
masks = masks.to(device)
outputs = model.forward(images)
The size of images and masks is (1, 1, 128, 128) when I print it out. This is because in my getitem_ in the Dataset object, i’m using image = torch.reshape(image, shape=(1, 128, 128, 128))
to reshape my image to the designated size. I tried changing that to image = torch.reshape(image, shape=(128, 128, 128)) and this doesn’t work either.
This is my train_loader:
train_loader = DataLoader(dataset=Dataset(partition['orig'], partition['segment']),
batch_size = batch_size, shuffle = True)
partition is just a dictionary that has original and segmentation mask images in it.