Find out specific module/layer names to restrict weights in PyTorch, MobileNetV3

Question

I want to restrict the range of a specific layer for the MobileNetV3. But I struggle with finding out the right module name, where to apply my defined constraints.

As an example: This is the first Inverted Residual block from the MobileNetV3, I only want to set constraints for the Conv2D layer, in the 0 ConvNormActivation part.

(1): InvertedResidual( (block): Sequential( (0): ConvNormActivation( (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=16, bias=False) (1): BatchNorm2d(16, eps=0.001, momentum=0.01, affine=True, track_running_stats=True) (2): ReLU(inplace=True)) (1): ConvNormActivation( (0): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False) (1): BatchNorm2d(16, eps=0.001, momentum=0.01, affine=True, track_running_stats=True)

So basically should say something like model._modules[‘conv2d layer in (block)sequential in (0)ConvNormActivation’].apply().

I tried to use:

for name, module in model.named_modules():
    print(name)

for the module name. But the module name “features.1.block.0.0” for the weights I consequently want to constrain (Weightname:"features.1.block.0.0.weight ") throws a “KeyError”.

model._modules[‘features.1.block.0.0’].apply(constraints)

I don’t want to constrain all weights within the model to the same boundaries, since the value range is quite different in some layers and therefore I would cut of some weights after training if I would not use different boundaries for different layers.

The constraints themselves are defined by:

class weightConstraint(object):
    def __init__(self):
        pass
    
    def __call__(self,module):
        if hasattr(module,'weight'):
            print("Entered")
            w=module.weight.data
            w=w.clamp(-0.3178,0.3921)
            module.weight.data=w

constraints=weightConstraint()

Applying the constraints to the whole model is not a problem, there I can just use:

model.apply(constraints)

But this constraints all weights within the model to the defined ones, I only want to apply them to certain layers for example. So my struggle has manly to do with finding the right names/modules.

For reference, I am injecting a BitFlip into the weights and therefore want to constrain the weights to the maximum and minimum values after training. So if a bitflip occurs the value does not go beyond the max and min value, and relies on the inherit fault resilience of DNN models for small value changes.

Any help would be greatly appreciated!