Hi!
I am currently writing a UNet implementation for medical image segmentation. However, I am struggling with memory problems, both when running with GPU and CPU. I’ve already read some posts and managed to solve the leaks by deleting intermediate tensors when doing the forward pass, but still my model requires around 16GB of RAM at some moments.
When only loading the data, this already requires 2.5GB of RAM. Then performing a forward pass with a batch-size of 5 pushes this to 16GB at most. You can find most of my code below, all tips on how to write memory efficient code or inspecting what parts require a lot of memory are really appreciated.
Dataset
class MedicalData(Dataset):
patches = []
labels = []
def __init__(self, dataset_location, transform=None):
self.transform = transform
for file in os.listdir(dataset_location):
filename = os.fsdecode(file)
#Load the training data from pickle files
if filename.startswith('patches'):
with open(dataset_location + filename, 'rb') as handle:
patches = pickle.load(handle)
self.patches.append(patches)
masks_filename = 'masks' + filename.split('patches')[1]
print(filename, masks_filename)
with open(dataset_location + masks_filename, 'rb') as handle:
labels = pickle.load(handle)
self.labels.append(labels)
#Flatten the lists into one list
self.patches = [item for sublist in self.patches for item in sublist]
self.labels = [item for sublist in self.labels for item in sublist]
self.patches = self.patches[:10]
self.labels = self.labels[:10]
assert (len(self.patches) == len(self.labels))
def __getitem__(self, index):
patch = self.patches[index]
patch = np.expand_dims(patch, axis=0)
if self.transform is not None:
patch = self.transform(patch)
# Convert patch and label to torch tensors
patch = torch.from_numpy(np.asarray(patch))
label = torch.from_numpy(np.asarray(self.labels[index]))
#Convert uint8 to float tensors
patch = patch.type(torch.FloatTensor)
label = label.type(torch.FloatTensor)
return patch, label
UNet
class Unet(nn.Module):
#input_channels is number of channels in input image
#num_filters is the amount of filters in the first conv layer
def __init__(self, input_channels, num_classes, num_filters, depth, padding=False):
super(Unet, self).__init__()
self.input_channels = input_channels
self.num_classes = num_classes
self.num_filters = num_filters
self.depth = depth
self.padding = padding
print("No of classes: %d \nNo of input channels: %d \nNo of filters first layer: %d \nDepth of the network: %d \nPadding: %d"
% (self.num_classes, self.input_channels, self.num_filters, self.depth, self.padding))
self.contracting_path = nn.ModuleList()
for i in range(depth):
input = self.input_channels if i == 0 else output
output = self.num_filters*(2**i)
self.contracting_path.append(DownConvBlock(input, output, padding))
self.upsampling_path = nn.ModuleList()
for i in range(depth-1):
input = output
output = input // 2
self.upsampling_path.append(UpConvBlock(input, output, padding))
self.last_layer = nn.Conv2d(output, num_classes, kernel_size=1)
def forward(self, x):
blocks = []
for i, down in enumerate(self.contracting_path):
x = down(x)
if i != len(self.contracting_path)-1:
blocks.append(x)
#x = F.avg_pool2d(x, 2)
x = F.max_pool2d(x, 2)
for i, up in enumerate(self.upsampling_path):
x = up(x, blocks[-i-1])
del blocks #Delete to fix memory leak
return self.last_layer(x)
class DownConvBlock(nn.Module):
def __init__(self, input_dim, output_dim, padding):
super(DownConvBlock, self).__init__()
layers = []
layers.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, padding=int(padding)))
layers.append(nn.BatchNorm2d(output_dim))
layers.append(nn.ReLU(inplace=True)) #Inplace is true is used to save memory
layers.append(nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=int(padding)))
layers.append(nn.BatchNorm2d(output_dim))
layers.append(nn.ReLU(inplace=True))
self.layers = nn.Sequential(*layers)
def forward(self, patch):
return self.layers(patch)
class UpConvBlock(nn.Module):
def __init__(self, input_dim, output_dim, padding, bilinear=False):
super(UpConvBlock, self).__init__()
if bilinear:
#self.upconv_layer = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.upconv_layer = nn.Sequential(nn.Upsample(mode='bilinear', scale_factor=2, align_corners=True), nn.Conv2d(input_dim, output_dim, kernel_size=1))
else:
self.upconv_layer = nn.ConvTranspose2d(input_dim, output_dim, kernel_size=2, stride=2)
self.conv_block = DownConvBlock(input_dim, output_dim, padding)
def forward(self, x, bridge):
up = self.upconv_layer(x)
crop1 = self.center_crop(bridge, up.shape[2:])
out = torch.cat([up, crop1], 1)
del up
del crop1
return self.conv_block(out)
