From 69bc3dd27b3af3955f52af333c72313f3997c6b6 Mon Sep 17 00:00:00 2001 From: CHIRANJEET1729DAS <126919390+CHIRANJEET1729DAS@users.noreply.github.com> Date: Mon, 9 Sep 2024 21:03:19 +0530 Subject: [PATCH] Update basicblock.py Changed conv layers --> lazy conv layers #Memory efficient code **refer comments for details** --- models/basicblock.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/models/basicblock.py b/models/basicblock.py index 12b8404b..bb99f106 100644 --- a/models/basicblock.py +++ b/models/basicblock.py @@ -58,13 +58,19 @@ def sequential(*args): # -------------------------------------------- # return nn.Sequantial of (Conv + BN + ReLU) # -------------------------------------------- -def conv(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bias=True, mode='CBR', negative_slope=0.2): +def conv(out_channels=64, kernel_size=3, stride=1, padding=1, bias=True, mode='CBR', negative_slope=0.2, in_channels=None): + # Changed convolution layers to lazy convolution layers. + # Lazy layers defer parameter initialization until input is passed through. + # This avoids the need to explicitly define input sizes, making the model more flexible. + # It also improves memory efficiency as memory is allocated only when the layers are used. L = [] for t in mode: if t == 'C': - L.append(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)) + L.append(nn.LazyConv2d(out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)) elif t == 'T': - L.append(nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)) + if in_channels is None: + raise ValueError("in_channels must be provided for ConvTranspose2d") + L.append(nn.LazyConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)) elif t == 'B': L.append(nn.BatchNorm2d(out_channels, momentum=0.9, eps=1e-04, affine=True)) elif t == 'I': @@ -94,7 +100,7 @@ def conv(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1, bi elif t == 'A': L.append(nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=0)) else: - raise NotImplementedError('Undefined type: '.format(t)) + raise NotImplementedError(f'Undefined type: {t}') return sequential(*L)