nvidia/C-RADIOv4-H

1# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9from logging import getLogger10import math11import os12from typing import Dict, List, Optional, Union, Tuple13from types import MethodType14 15import torch16from torch import nn17from torch.nn import functional as F18from torch.nn.utils import parametrize19 20 21# For now, don't do anything22class DAMP(nn.Identity):23    def __init__(self, std: float):24        super().__init__()25        self.std = std26 27 28def enable_damp(model: nn.Module, std: float):29    if isinstance(model, (list, tuple)):30        for m in model:31            enable_damp(m, std)32        return33 34    for name, module in model.named_modules():35        if isinstance(module, nn.Linear):36            parametrize.register_parametrization(module, 'weight', DAMP(std))37 38 39def configure_damp_from_args(model: nn.Module, args):40    damp = getattr(args, 'damp', None)41    if damp:42        enable_damp(model, damp)43