From 3D datas to 2D slices

oguzcan05: <p>Hi there I have faced a problem and then I have read some forums but I couldn`t get the correct answer. I would like to ask again, in Dataset part I would like to read my 3D data which has first 2 dimension same shape ([512,512,600] or [512,512,550]) and I would like to create like 600 slices or 550 slices which has just [512,512] dimensions as 2d array.</p>
<h1><a name="heading-1" class="anchor" href="#heading-1"></a></h1>
<pre><code class="lang-auto">class TrainSet(Dataset):
    def __init__(self, image_dir, mask_dir):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.images = os.listdir(image_dir)
        self.masks = os.listdir(mask_dir)

  

    def __getitem__(self, index):
        img_path = os.path.join(self.image_dir, self.images[index])
        mask_path = os.path.join(self.mask_dir, self.masks[index])
        image = nibabel.load(img_path)
        img = image.get_fdata()
        mask, _ = nrrd.read(mask_path)
        
        for k in range(img.shape[2]):
            img = img[:, :, k]
            mask = mask[:, :, k]
            
            return img, mask

    def __len__(self):
        return len(self.images)
</code></pre>
<p>I have tried something like that but as you guess it couldn<code>t apply for whole images. And I didn</code>t prefer to make list by slice by actually because it causes other problems. At that point can you suggest me some solutions please? How can I create my 2D dataset from 3D arrays? Thank you very much in advance~~</p>

ptrblck: <p>I don’t know how these slices should be used, i.e. if you want to push each slice separately to the model or if you would like to create a batch of slices from different files.<br>
In the former case it might be easier to just return all slices from the current file and iterate them in the training loop.<br>
In the latter case, a custom sampler (e.g. using two indices for the file and slice index) might work.</p>

oguzcan05: <p>Thank you very much for fast reply~<br>
I think both idea looks okay but the first one is muc more close. From your suggestion I understood that I should return 3d arrays and seperate on the training loop… I am sorry I couldn`t understand clearly how can I iterate them in training loop to 2D arrays?</p>

ptrblck: <p>Yes, you could directly return the 3D tensor, which would have the shape <code>[batch_size=1, 512, 512, 600]</code> inside the training loop.<br>
Afterwards, you could index each slice via:</p>
<pre><code class="lang-python">for data, target in loader:
    # data will have the shape [1, 512, 512, 600]
    for idx in data.size(3):
        x = data[:, :, :, i]
        # x will have the shape [1, 512, 512]
</code></pre>
<p>and could then pass <code>x</code> to the model.</p>

oguzcan05: <p>I would like to try that way but in advance why I am having this problem: RuntimeError: stack expects each tensor to be equal size, but got [512, 512, 533] at entry 0 and [512, 512, 661] at entry 1<br>
I didn`t change the code I just deleted the for part…</p>
<pre><code class="lang-auto">class TrainSet(Dataset):
    def __init__(self, image_dir, mask_dir):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.images = os.listdir(image_dir)
        self.masks = os.listdir(mask_dir)

    def __getitem__(self, index):
        img_path = os.path.join(self.image_dir, self.images[index])
        mask_path = os.path.join(self.mask_dir, self.masks[index])
        image = nibabel.load(img_path)
        img = image.get_fdata()
        mask, _ = nrrd.read(mask_path)



        return img, mask

    def __len__(self):
        return len(self.images)
</code></pre>

ptrblck: <p>This error is raised in the <code>collate_fn</code> which tries to create a batch of multiple samples.<br>
Since each sample has a different shape, <code>torch.stack</code> will fail.<br>
If you want to iterate each slice of a single sample, you could thus set <code>batch_size=1</code> and it should work.</p>

oguzcan05: <p>I am really thankful for fast replying I really need to finish this code so I am also sorry for asking so question~~<br>
After I finished my dataset part I was started to train part like that basically:</p>
<pre><code class="lang-auto">device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = UNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)


for epoch in range(epochs):
    for i, (images, masks) in enumerate(train_loader, 1):
        print(images.size())
        for idx in range(images.size(3)):
            images = images[:, :, :, idx].unsqueeze(1).to(device)
            masks = masks[:, :, :, idx].unsqueeze(1).to(device)
            images = images.type(torch.FloatTensor)
            masks = images.type(torch.FloatTensor)
            print(images.shape)
            print(masks.shape)
            outputs = model(images)
            loss = criterion(outputs, masks)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i) % 1 == 0:
                print(f"Epoch [{epoch + 1}/{epochs}], Step [{i}/{total_steps}], Loss: {loss.item():4f}")
</code></pre>
<p>But I couldn`t understand why it was connected to CrossEntropyLoss because it is giving this error from loss part : RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4</p>

ptrblck: <p>I guess you are working on a multi-class segmentation?<br>
If so, the model output should have the shape <code>[batch_size, nb_classes, height, width]</code>, while the target should have the shape <code>[batch_size, height, width]</code> and container class indices in the range <code>[0, nb_classes-1]</code> as shown here:</p>
<pre><code class="lang-python">batch_size, nb_classes, height, width = 2, 3, 4, 4
output = torch.randn(batch_size, nb_classes, height, width)
target = torch.randint(0, nb_classes, (batch_size, height, width))

criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)
</code></pre>
<p>Based on your error message I guess your target has 4 dimensions, which is wrong. If you are using a one-hot encoded target, use <code>target = torch.argmax(target, dim=1)</code> to create the expected shape (assuming the class dimension is in <code>dim1</code>).</p>