def set_requires_grad(model: Model, requires_grad: Union[bool, Dict[str,
bool]]):
"""Sets the ``requires_grad`` value for all model parameters.
Example::
>>> model = SimpleModel()
>>> set_requires_grad(model, requires_grad=True)
>>> # or
>>> model = SimpleModel()
>>> set_requires_grad(model, requires_grad={""})
Args:
model (torch.nn.Module): model
requires_grad (Union[bool, Dict[str, bool]]): value
"""
if isinstance(requires_grad, dict):
for name, param in model.named_parameters():
assert (name in requires_grad
), f"Parameter `{name}` does not exist in requires_grad"
param.requires_grad = requires_grad[name]
else:
requires_grad = bool(requires_grad)
for param in model.parameters():
param.requires_grad = requires_grad
def set_requires_grad(model: Model, requires_grad: bool):
"""Sets the ``requires_grad`` value for all model parameters.
Example::
>>> model = SimpleModel()
>>> set_requires_grad(model, requires_grad=True)
Args:
model (torch.nn.Module): model
requires_grad (bool): value
"""
requires_grad = bool(requires_grad)
for param in model.parameters():
param.requires_grad = requires_grad