nvidia/C-RADIOv2-VLM-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.8from typing import Optional9 10import torch11from torch import nn12 13 14class ClsToken(nn.Module):15    def __init__(self, ndim: int,16                 num_tokens: int = 1,17                 enabled: bool = True,18                 register_multiple: Optional[int] = None,19                 num_registers: Optional[int] = None,20    ):21        super().__init__()22 23        self.ndim = ndim24        self.enabled = enabled25        self.num_registers = 026        self.num_tokens = num_tokens27        if enabled:28            if num_registers:29                self.num_registers = num_registers30            elif register_multiple:31                self.num_registers = register_multiple - (num_tokens % register_multiple)32 33            scale = ndim ** -0.534            self.token = nn.Parameter(torch.randn(num_tokens + self.num_registers, ndim) * scale)35        else:36            self.token = None37 38        self.num_patches = self.num_tokens + self.num_registers39 40    def disable(self):41        self.token = None42        self.enabled = False43 44    def forward(self, x: torch.Tensor):45        if self.token is None:46            return x47 48        token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)49        x = torch.cat([50            token,51            x,52        ], dim=1)53 54        return x55 56    def no_weight_decay(self):57        return [58            'token',59        ]60