How to call specific layer index during training?

username: lukman
cooked: <p>I am tying to modify ResNet34 for Task Incremental learning as part of my academic work. It done by replacing every BatchNorm2d locations with Sequential block of multiple BatchNorm2d layers. Let say have a Sequential block of 4 BN layers at each location of BN in ResNet34. For training, I will do separate training on each BN blocks (here, 4 trainig) while using all other CNN layers (with task specific masks) during all training. Issue is, how can I call or set specific BN sublayer (like as index) in the Sequential block prior to its training inorder to use that BN layer. Ex; set model to use BN sublayer 2 (2 index) at every Sequential() block for 2nd training.</p>
<p>I am able to modify model like what I mentioned. But not able to train likey. What I noticed is; I have set an attribute self.bn_nr initialized with 0, so it able to train for 1st index (or 0), and for subsequent training model.bn_nr is getting updated. but not find any scope to update this attribute for MultiBatchNorm() forward() method for next training. Is there anyway to get this updation on self.bn_nr within forward() of MultiBatchNorm(). Can anyone help me out to solve this?</p>
<p>Here I skipped Incremental learning part and trying to summarizing my problem with least code because it might be redundant with the question I asked.</p>
<pre><code>class BaseModel(nn.Module, metaclass=abc.ABCMeta):
    """Abstract module to add CL capabilities to multitask classifier"""

    def __init__(self):
        super().__init__()
        #self.scenario = None
        self.num_bn_layers = 0
        self.bn_nr = 0

    @abc.abstractmethod
    def forward(self, x):
        pass

class MultiBatchNorm(BaseModel):
    def __init__(self, num_features=None, num_bn_layers=4):
        super().__init__()
        self.batchnorms = nn.ModuleList([nn.BatchNorm2d(num_features) for _ in range(num_bn_layers)])
        self.num_bn_layers = num_bn_layers

    def forward(self, x):
        x = self.batchnorms[self.bn_nr](x)
        return x

class ResidualBlock(BaseModel):
    def __init__(self, in_channels, out_channels, num_bn_layers, stride = 1, downsample = None):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Sequential(
                        nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride = stride, padding = 1),
                        MultiBatchNorm(out_channels, num_bn_layers),
                        nn.ReLU())
        self.conv2 = nn.Sequential(
                        nn.Conv2d(out_channels, out_channels, kernel_size = 3, stride = 1, padding = 1),
                        MultiBatchNorm(out_channels, num_bn_layers))
        self.downsample = downsample
        self.relu = nn.ReLU()
        self.out_channels = out_channels

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.conv2(out)
        if self.downsample:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out

class ResNet(BaseModel):
    def __init__(self, block, layers, num_classes=10, classes=10, num_bn_layers=2):
        super(ResNet, self).__init__()
        self.inplanes = 64
        self.conv1 = nn.Sequential(
                        nn.Conv2d(3, 64, kernel_size = 7, stride = 2, padding = 3),
                        MultiBatchNorm(64, num_bn_layers),
                        nn.ReLU())
        self.maxpool = nn.MaxPool2d(kernel_size = 3, stride = 2, padding = 1)
        self.layer0 = self._make_layer(block, 64, layers[0], num_bn_layers, stride = 1)
        self.layer1 = self._make_layer(block, 128, layers[1], num_bn_layers, stride = 2)
        self.layer2 = self._make_layer(block, 256, layers[2], num_bn_layers, stride = 2)
        self.layer3 = self._make_layer(block, 512, layers[3], num_bn_layers, stride = 2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512, classes)

        self.num_bn_layers = num_bn_layers

    def _make_layer(self, block, planes, blocks, num_bn_layers, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes:

            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes, kernel_size=1, stride=stride),
                MultiBatchNorm(planes, num_bn_layers),
            )
        layers = []
        layers.append(block(self.inplanes, planes, num_bn_layers, stride, downsample))
        self.inplanes = planes
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, num_bn_layers))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = nn.functional.interpolate(x, (224, 224))
        x = self.conv1(x)
        x = self.maxpool(x)
        x = self.layer0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x
        

num_bn_layers = 4
model = ResNet(ResidualBlock,
               layers=[3, 4, 6, 3],
               num_bn_layers=num_bn_layers,
               ).to(device)
</code></pre>
<h2><a name="training-pseudo-code-1" class="anchor" href="#training-pseudo-code-1"></a>training (pseudo code)</h2>
<pre><code>epochs = 10
for task_id in range(num_bn_layers):
  model.bn_nr = task_id
  for epoch in range(epochs):
    for b, data in enumerate(train_loader):
      X, y = data
      loss = model.train_epoch(X)
</code></pre>