33 lines
1003 B
Python
33 lines
1003 B
Python
"""Contains utility functionality to modify torch modules.
|
|
|
|
For licensing see accompanying LICENSE file.
|
|
Copyright (C) 2025 Apple Inc. All Rights Reserved.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from torch import nn
|
|
|
|
NORM_LAYER_TYPES = tuple(module_type for name, module_type in nn.__dict__.items() if "Norm" in name)
|
|
BATCH_NORM_LAYER_TYPES = tuple(
|
|
module_type for name, module_type in nn.__dict__.items() if "BatchNorm" in name
|
|
)
|
|
|
|
|
|
def freeze_norm_layer(module: nn.Module) -> nn.Module:
|
|
"""Freeze all normalization layers."""
|
|
|
|
def set_module_eval_mode(module: nn.Module, _: Any) -> None:
|
|
module.eval()
|
|
|
|
for submodule in module.modules():
|
|
if isinstance(submodule, NORM_LAYER_TYPES):
|
|
submodule.requires_grad_(False)
|
|
# This is to ensure that batch norm layers are always called
|
|
# with the precomputed running statistics.
|
|
submodule.register_forward_pre_hook(set_module_eval_mode)
|
|
|
|
return module
|