mirror of https://github.com/razor-ai/soup.git
74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
"""Freeze training: freeze bottom N layers of a model for parameter-efficient training."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from typing import Any, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _detect_num_layers(model: Any) -> int:
|
|
"""Detect total number of transformer layers from model parameter names.
|
|
|
|
Looks for patterns like 'model.layers.N.' or 'transformer.h.N.' and
|
|
returns max(N) + 1.
|
|
"""
|
|
max_layer = -1
|
|
pattern = re.compile(r"(?:layers|h)\.(\d+)\.")
|
|
for name, _ in model.named_parameters():
|
|
match = pattern.search(name)
|
|
if match:
|
|
layer_idx = int(match.group(1))
|
|
if layer_idx > max_layer:
|
|
max_layer = layer_idx
|
|
return max_layer + 1 if max_layer >= 0 else 0
|
|
|
|
|
|
def freeze_model_layers(
|
|
model: Any,
|
|
freeze_layers: Optional[int] = None,
|
|
freeze_ratio: Optional[float] = None,
|
|
) -> int:
|
|
"""Freeze the bottom layers of a model.
|
|
|
|
Args:
|
|
model: A PyTorch model with named_parameters().
|
|
freeze_layers: Freeze the first N layers. Takes priority over freeze_ratio.
|
|
freeze_ratio: Freeze this fraction of layers (e.g. 0.75 = 75% from bottom).
|
|
|
|
Returns:
|
|
Number of parameters frozen.
|
|
"""
|
|
if freeze_layers is None and freeze_ratio is None:
|
|
return 0
|
|
|
|
total_layers = _detect_num_layers(model)
|
|
if total_layers == 0:
|
|
logger.warning(
|
|
"freeze_model_layers: could not detect numbered layers in model "
|
|
"parameter names. Freezing has no effect. Check that your model "
|
|
"uses 'layers.N.' or 'h.N.' naming."
|
|
)
|
|
return 0
|
|
|
|
# Determine cutoff
|
|
if freeze_layers is not None:
|
|
cutoff = min(freeze_layers, total_layers)
|
|
else:
|
|
cutoff = int(total_layers * freeze_ratio)
|
|
|
|
# Freeze parameters in layers below cutoff
|
|
frozen_count = 0
|
|
pattern = re.compile(r"(?:layers|h)\.(\d+)\.")
|
|
for name, param in model.named_parameters():
|
|
match = pattern.search(name)
|
|
if match:
|
|
layer_idx = int(match.group(1))
|
|
if layer_idx < cutoff:
|
|
param.requires_grad = False
|
|
frozen_count += 1
|
|
|
|
return frozen_count
|