Scribd
Upload
13 views3 pages

Hrhrttgggegerge

The document provides guidance on using fast.ai's ImageDataLoaders with a specified batch size for training models on limited GPU memory. It includes a code example for setting up a UNet model and implementing checkpointing to manage memory efficiently during training. The document emphasizes the importance of adjusting the batch size and utilizing memory management techniques in deep learning workflows.

Uploaded by

Hoàng Văn Nguyên
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as DOCX, PDF, TXT or read online on Scribd
13 views3 pages

Hrhrttgggegerge

The document provides guidance on using fast.ai's ImageDataLoaders with a specified batch size for training models on limited GPU memory. It includes a code example for setting up a UNet model and implementing checkpointing to manage memory efficiently during training. The document emphasizes the importance of adjusting the batch size and utilizing memory management techniques in deep learning workflows.

Uploaded by

Hoàng Văn Nguyên
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as DOCX, PDF, TXT or read online on Scribd
You are on page 1/ 3

If someone arrives here because of fast.

ai, the batch size of a loader


such as ImageDataLoaders can be controlled via bs=N where N is the size of
the batch.

My dedicated GPU is limited to 2GB of memory, using bs=8 in the


following example worked in my situation:

from fastai.vision.all import *

path = untar_data(URLs.PETS)/'images'

def is_cat(x): return x[0].isupper()

dls = ImageDataLoaders.from_name_func(

path, get_image_files(path), valid_pct=0.2, seed=42,

label_func=is_cat, item_tfms=Resize(244), num_workers=0, bs=)

learn = cnn_learner(dls, resnet34, metrics=error_rate)

learn.fine_tune(1)

You can try something like this before your training loop

model = UNET( n_channels, n_classes)

for i in range(epochs):

torch.cuda.empty_cache()

model.use_checkpointing() # Or directly you can do 'model =


torch.utils.checkpoint(model)'

your_train_function()

where your model definition should be like below block

class UNet(nn.Module):

def __init__(self, n_channels, n_classes, bilinear=False):

super(UNet, self).__init__()
self.n_channels = n_channels

self.n_classes = n_classes

self.bilinear = bilinear

self.inc = (DoubleConv(n_channels, 64))

self.down1 = (Down(64, 128))

self.down2 = (Down(128, 256))

self.down3 = (Down(256, 512))

factor = 2 if bilinear else 1

self.down4 = (Down(512, 1024 // factor))

self.up1 = (Up(1024, 512 // factor, bilinear))

self.up2 = (Up(512, 256 // factor, bilinear))

self.up3 = (Up(256, 128 // factor, bilinear))

self.up4 = (Up(128, 64, bilinear))

self.outc = (OutConv(64, n_classes))

def forward(self, x):

x1 = self.inc(x)

x2 = self.down1(x1)

x3 = self.down2(x2)

x4 = self.down3(x3)

x5 = self.down4(x4)

x = self.up1(x5, x4)

x = self.up2(x, x3)

x = self.up3(x, x2)

x = self.up4(x, x1)

logits = self.outc(x)
return logits

def use_checkpointing(self):

self.inc = torch.utils.checkpoint(self.inc)

self.down1 = torch.utils.checkpoint(self.down1)

self.down2 = torch.utils.checkpoint(self.down2)

self.down3 = torch.utils.checkpoint(self.down3)

self.down4 = torch.utils.checkpoint(self.down4)

self.up1 = torch.utils.checkpoint(self.up1)

self.up2 = torch.utils.checkpoint(self.up2)

self.up3 = torch.utils.checkpoint(self.up3)

self.up4 = torch.utils.checkpoint(self.up4)

self.outc = torch.utils.checkpoint(self.outc)

more details in the following link

You might also like