About large datasize, 3D data and patches

banikr: <p>Hello All,<br>
I am working on 3D data of 114 images each of dimensions <code>[180x256x256]</code>. Since such a large image can not be fed directly to the network, I am using overlapping patches of size <code>[64x64x64]</code>. Now there are around 22,000 patches in total for 114 images. which can not be loaded into the <code>Dataloader</code> as <code>cuda</code> memory runs out.<br>
Is there a way to iterate the loading of one image at a time and run patches from the image and go to the next image?</p>
<p>N.B. for each image there is a target 3D mask with 9 different labels.</p>

ptrblck: <p>Often you would not load the data directly to the GPU in the <code>Dataset.__getitem__</code> or in any <code>DataLoader</code> method, but you would push it to the device in the training loop.<br>
Assuming your host system has enough memory to load (a batch of) images, you could perform the patch creation in the <code>Dataset</code>.<br>
Depending on your use case, you could e.g. specify the batch size to be <code>N</code> and slice the single image into <code>M</code> patches. The returned batch would then be <code>N*M</code> of the smaller windows.</p>

banikr: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a>,<br>
I wrote some of the <code>Dataset</code> like below:</p>
<pre><code class="lang-auto">class SegSet(data.Dataset):
    def __init__(self, subdict, num_labels):  
        """
        :param subdict: a dictionary of subject and label paths
        :param num_labels: number of segmentation labels (9)
        """
        self.subdict = subdict
        self.img_subs = subdict['img_subs']
        self.img_files = subdict['img_files']
        if checkKey(subdict, 'seg_subs'): # This function checks if segmentation label available or not
            self.seg_subs = subdict['seg_subs']
            self.seg_files = subdict['seg_files']
        else:
            self.seg_subs = None
            self.seg_files = None
        self.num_labels = num_labels

    def __len__(self):
        return len(self.img_subs)

    def __getitem__(self, index):
        num_labels = self.num_labels
        sub_name = self.img_subs[index]
        img_file = self.img_files[index]
        img_3d = nib.load(img_file)
        img = img_3d.get_data()
        img = (img - img.min())/(img.max()-img.min())
        img = img*255.0
        seg_file = self.seg_files[index]
        seg_3d = nib.load(seg_file)
        seg = seg_3d.get_data()
        imgp, segp = generate_patch_32_3(img, seg)
        for i in range(1,num_labels):
            for j in range(len(imgp)):
                seg_one = segp == labels[i] #labels = labels number list e.g. [0, 1, 2 ,10, 56 ...]
                segp[j, i, :, :, :] = seg_one[0:segp.shape[0], 0:segp.shape[1], 0:segp.shape[2]]
                segp[j, 0, :, :, :] = segp[j, 0, :, :, :] - segp[j, i, :, :, :]
                # print("Here")
        imgp = imgp.astype('float32')
        segp = segp.astype('float32')
        return imgp, segp, sub_name
</code></pre>
<p>The <code>generate_patch_32_3</code> function simply generates 3D patches from 180x256x256 img and segmentation img to 192x16x64x64 and 192x8x32x32 correspondingly.</p>
<p>The problem is after having paired source and target from <code>SegSet</code> I don’t understand how to load them to the 'Dataloader<code>and also keep loading the next image for</code>SegSet`.</p>
<pre><code class="lang-auto">    train_set = SegSet(train_dict, num_labels=9)
    print(len(train_set))
    x, y, z = next(iter(train_set))
    print(x.shape, '\n', y.shape, '\n', z)
</code></pre>
<p>The output is:<br>
<code>2</code><br>
<code>(192, 1, 16, 64, 64)</code><br>
<code>(192, 9, 8, 32, 32) </code><br>
<code>001_MR2std</code></p>
<p>So the output of the <code>Dataset</code> class is multiple (192) source and targets. When I pass it to <code>Dataloader</code> :</p>
<pre><code class="lang-auto">train_loader = data.DataLoader(train_set, batch_size=16, shuffle=False, num_workers=1)
print(len(train_loader))
x, y, z = next(iter(train_loader))
print(x.shape, '\n', y.shape, '\n', z)
</code></pre>
<p>The output is:</p>
<p><code> 1</code><br>
<code>torch.Size([2, 192, 1, 16, 64, 64]) </code><br>
<code>torch.Size([2, 192, 9, 8, 32, 32]) </code><br>
<code>('001_MR2std', '002_MR2std')</code></p>
<p>It doesn’t change even if I change the <code>batch_size</code> to 8 or 80.</p>
<p>So if I feed the x and y to the network, is it taking the whole 192 patches at once? what is the significance of batches here?</p>

ptrblck: <p>I’m not sure what exactly is causing the issue and cannot easily debug it as I don’t have the files.<br>
However, here is a small example of what I had in mind:</p>
<pre><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.randn(10, 3, 24, 24)
    
    def __getitem__(self, index):
        # get current sample
        x = self.data[index]
        # create patches
        x = x.unfold(1, 12, 12).unfold(2, 12, 12)
        #flatten patches
        x = x.contiguous().view(x.size(0), -1, x.size(3), x.size(4))
        x = x.permute(1, 0, 2, 3)
        return x
        
    def __len__(self):
        return len(self.data)

dataset = MyDataset()
loader = DataLoader(dataset, batch_size=2)
x = next(iter(loader))
print(x.shape)
&gt; torch.Size([2, 4, 3, 12, 12])
</code></pre>
<p>Inside the <code>Dataset.__getitem__</code> you are creating 4 patches of the shape <code>12x12</code> and returning these 4 patches in a shape of <code>[patches=4, channels=3, height=12, width=12]</code>.<br>
The <code>DataLoader</code> uses the <code>batch_size</code> to specify how often the <code>__getitem__</code> is called and stacks the returned tensors in a newly created <code>dim0</code>.<br>
For a batch size of 2, the returned batch will have thus the shape <code>[batch_size=2, patches=4, 3, 12, 12]</code> and you could flatten the patches into the batch dimension in order to feed it into the model.</p>

banikr: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a>,<br>
I think I have understood the issue. Since I fed only patches from 2 images, thus the batch never changes beyond 2. Sorry for the confusion.</p>
<p>Now the questions remain,<br>
I am giving <code>Dataset</code> one image and one segmentation path in <code>index</code>, which loads one image[170,256,256] and segmentation[same size] file from directory. As you suggested, I am generating patches inside <code>__getitem__</code><br>
after patch operation,<br>
img: [<strong>192</strong>, 16, 64, 64]<br>
seg: [<strong>192</strong>, 8, 32, 32]<br>
finally I create segmentation channels for 9 labels<br>
img: [192, <strong>1</strong> , 16, 64, 64]<br>
seg: [192, <strong>9</strong>, 8, 32, 32]</p>
<p>This way I can directly load data from the directories and need not save anything.</p>
<p>Is there a way to iterate 192 patches as<br>
src: <code>[batch_size, 1, 16, 64, 64]</code><br>
trgt:<code>[batch_size, 9, 8, 32, 32]</code> and feed to <code>Dataloader</code> and also iterate 3D images in <code>Dataset</code>?</p>
<p>Now to do the same I had to save all the patches to an hdf5(around 8gb) file as and index one by one.</p>
<pre><code class="lang-auto">    g = h5py.File(hpath[0], 'r')
    print(g.keys())
    img = np.array(g['MR'])    # (114, 192, 16, 64, 64)
    seg = np.array(g['Mask'])   # (114, 192, 8, 32, 32)


    train_img = img.reshape(-1, 16, 64, 64)
    train_seg = seg.reshape(-1, 8, 32, 32)

    trainSet = SimpleDataset(train_img, train_seg)
    x, y = next(iter(trainSet))
    print(x.shape, '\n', y.shape)
    # (1, 1, 16, 64, 64)
    # (1, 9, 8, 32, 32)
</code></pre>
<p>Is there a more efficient way to do that?<br>
I hope I make sense and convey it clearly <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>

ptrblck: <aside class="quote no-group" data-username="banikr" data-post="5" data-topic="112630">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/banikr/48/54353_2.png" class="avatar"> banikr:</div>
<blockquote>
<p>Is there a way to iterate 192 patches as<br>
src: <code>[batch_size, 1, 16, 64, 64]</code><br>
trgt:<code>[batch_size, 9, 8, 32, 32]</code> and feed to <code>Dataloader</code> and also iterate 3D images in <code>Dataset</code>?</p>
</blockquote>
</aside>
<p>I’m not sure I understand the issue correctly.<br>
The <code>Dataset</code> approach would create the patches and the <code>DataLoader</code> would then return tensors with <code>batch_size*num_patches</code> in <code>dim0</code>, which you could feed directly into the model.<br>
Why would you have to store the patches again?</p>

banikr: <p>so if I generate patches inside <code>__getitem__</code> it gives <code>[192, 1, 64, 64, 64]</code> img from <code>Dataset</code>. If I pass it through <code>Dataloader</code> with <code>batch_size</code> 2, finally source is <code>[2, 192, 1, 1, 64, 64, 64]</code><br>
Even if I reshape it to <code>[384, 1, 64, 64, 64]</code></p>
<p>How would I feed it to the model?</p>
<p>[N.B. input to the <code>Dataset</code> is actually location to each 3D image paths.]</p>

ptrblck: <p>You would feed it in the “flattened” shape to the model, i.e. <code>[batch_size*num_patches=384, channels=1, ...]</code>.<br>
This would increase the “new batch size” by a factor of <code>num_patches</code>, but would still work.<br>
Let me know, if I’m still misunderstanding the issue or what kind of errors you are getting.</p>

banikr: <p>In that case, GPU runs out of memory:</p>
<p><code>&gt;&gt; RuntimeError: CUDA out of memory. Tried to allocate 3.00 GiB (GPU 0; 10.89 GiB total capacity; 4.70 GiB already allocated; 2.16 GiB free; 4.70 GiB reserved in total by PyTorch)</code></p>

ptrblck: <p>In that case you could lower the batch size to 1 and check, if it’s still running out of memory.<br>
If that’s the case, keep the batch size at 1, split the batch in the <code>DataLoader</code> loop in <code>dim0</code> and loop over smaller input tensors.<br>
Alternatively, you could also try to create the patches from each image in several steps in the <code>__getitem__</code>, but I think this would be much more complicated, as you might need a custom sampler for it etc.</p>

banikr: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
I think the looping works:</p>
<pre><code class="lang-auto">    train_loader = DataLoader(train_set, batch_size=1, shuffle=True, num_workers=8)
    print(len(train_loader))
    # x, y, z = next(iter(train_loader))
    # print(x.shape, '\n', y.shape, '\n', z) #torch.Size([1, 192, 1, 16, 64, 64])
                                           #torch.Size([1, 192, 9, 8, 32, 32])
    model = UNet3D(in_channel=1, n_classes=9).to(device)
    bSize = 4
    # the loop below would be within epochs
    for idx, sample in enumerate(train_loader):
        print('idx: {}'.format(idx), 'sub: {}'.format(sample[2]))
        for i in range(0, sample[0].shape[1], bSize):
            img = sample[0][0, i:i+bSize, 0, :, :, :] # creating batches
            img = img.unsqueeze(1)
            # print(img.shape)
            seg = sample[1][0, i:i+bSize, 0:9, :, :, :]
            # print(seg.shape)
            img, seg = img.to(device), seg.to(device)
            out = model(img)
            # print('Out: {}'.format(out.shape))
</code></pre>
<p>The same code also iterates the main subject images, creates the patches from each inside <code>Dataset</code>, takes samples from <code>Dataloader</code>, and chunks from sample fed to GPU/model at a time.<br>
Thanks.</p>
<p>Is there a better way to check that correct data are sent to the model, apart from the tensor shape?</p>
<p>Just out of curiosity, are there Pytorch functionality for that? Some of the posts I have seen has words like <code>Torchnet</code>, <code>ConcatDataset</code>, <code>ChunkSize</code> and things like that?</p>

ptrblck: <aside class="quote no-group" data-username="banikr" data-post="11" data-topic="112630">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/banikr/48/54353_2.png" class="avatar"> banikr:</div>
<blockquote>
<p>Is there a better way to check that correct data are sent to the model, apart from the tensor shape?</p>
</blockquote>
</aside>
<p>You could create pre-defined tensors as the input to the <code>Dataset</code> and verify the values in the training loop.<br>
E.g. create 10 samples, where each image contains a specific scalar value (you could use <code>torch.arange</code>, <code>unsqueeze</code> the needed dimensions, and <code>expand</code> the tensor). Before passing the chunks to the model, print the data and make sure the samples look correct and are not e.g. interleaved.</p>
<aside class="quote no-group" data-username="banikr" data-post="11" data-topic="112630">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/banikr/48/54353_2.png" class="avatar"> banikr:</div>
<blockquote>
<p>Just out of curiosity, are there Pytorch functionality for that? Some of the posts I have seen has words like <code>Torchnet</code>, <code>ConcatDataset</code>, <code>ChunkSize</code> and things like that?</p>
</blockquote>
</aside>
<p><code>Torchnet</code> was a higher-level API for Torch7 (LuaTorch), if I’m not mistaken. <code>ConcatDataset</code> is used to concatenate different datasets to a single one (which wouldn’t help here), and I don’t know what <code>ChunkSize</code> refers to.<br>
That being said, I don’t think this functionality is added to PyTorch core modules, but might be available in a high-level API, such as <code>Ignite</code>, <code>Lightning</code>, or <code>Catalyst</code>.</p>

fepegar: <p>Hey, <a class="mention" href="/u/banikr">@banikr</a>. You might be interested in using TorchIO for your problem. I recommend you take a look at <a href="https://torchio.readthedocs.io/data/patch_based.html" class="inline-onebox" rel="noopener nofollow ugc">Patch-based pipelines - TorchIO 0.18.29 documentation</a>.</p>