Change specific named variables in all Modules

1

image
image875×440 6.7 KB

Hello, I have some problems changing pretrained models (in my case Yolov5).
Model is built by selections of modules and numbers of repetition.
I want to change modules like above, trying to mask specific spatial regions of input features per image inference.
However, I don’t know how to change all module variables named mask per infernece.

How can I check and replace mask values in all modules?
This is the logic flow I want:

> for module in all_Modules:
>    if mask in module variables:
>        mask = specific_values_I_want
2

You could iterate the modules of the parent nn.Module, check for the mask attribute via hasattr and assign e.g. a new nn.Parameter to it as seen here:

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.mask = None
        

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(1, 1)
        self.mod1 = MyModule()
        self.lin2 = nn.Linear(1, 1)
        self.mod2 = MyModule()


model = MyModel()
print(dict(model.named_parameters()))
> {'lin1.weight': Parameter containing:
  tensor([[-0.3372]], requires_grad=True), 'lin1.bias': Parameter containing:
  tensor([0.7264], requires_grad=True), 'lin2.weight': Parameter containing:
  tensor([[0.0947]], requires_grad=True), 'lin2.bias': Parameter containing:
  tensor([0.2308], requires_grad=True)}

for name, module in model.named_modules():
    if hasattr(module, 'mask'):
        print('{} has mask attribute'.format(name))
        module.mask = nn.Parameter(torch.randn(1))
> mod1 has mask attribute
  mod2 has mask attribute

print(dict(model.named_parameters()))
> {'lin1.weight': Parameter containing:
  tensor([[-0.3372]], requires_grad=True), 'lin1.bias': Parameter containing:
  tensor([0.7264], requires_grad=True), 'mod1.mask': Parameter containing:
  tensor([-0.6128], requires_grad=True), 'lin2.weight': Parameter containing:
  tensor([[0.0947]], requires_grad=True), 'lin2.bias': Parameter containing:
  tensor([0.2308], requires_grad=True), 'mod2.mask': Parameter containing:
  tensor([-0.1601], requires_grad=True)}
1 Like