I want to set submodule as no_grad for reduce memory.
Specifically, I consider the following example
import timm
model = timm.create_model("convnext_small_in22k")
stem = model.get_submodule("stem")
model.__setattr__("stem", torch.no_grad()(stem))
TypeError: cannot assign 'torch.nn.modules.container._DecoratorContextManager.__call__.<locals>.decorate_context' as child module 'stem' (torch.nn.Module or None expected)
Please tell me how to modify submodel to have no_grad attribute
(requires_grad = False is not a solution)
+ I don’t want to use a wrapper class that runs no_grad forward.
+ After applying no_grad , model.named_parameters() should return the same name.
+ The current makeshift is to use a wrapper class.
class no_grad_wrapper(nn.Module):
def __init__(self, module):
super().__init__()
self.M = module
@torch.no_grad()
def forward(self, x):
return self.M(x)
But this is not good because it creates a letter M
in named_parameters that doesn’t exist
How to handle it without creating a variable, or is it possible to apply no_grad statically?