56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
"""
|
|
-----------------------------------------------------------------------------
|
|
Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
and proprietary rights in and to this software, related documentation
|
|
and any modifications thereto. Any use, reproduction, disclosure or
|
|
distribution of this software and related documentation without an express
|
|
license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
-----------------------------------------------------------------------------
|
|
"""
|
|
|
|
from typing import Literal, Optional, Tuple
|
|
|
|
import attrs
|
|
|
|
|
|
@attrs.define(slots=False)
|
|
class ModelConfig:
|
|
# input
|
|
use_salient_point: bool = True
|
|
|
|
# random cutoff during training
|
|
cutoff_fps_point: Tuple[int, ...] = (256, 512, 512, 512, 1024, 1024, 2048)
|
|
cutoff_fps_salient_point: Tuple[int, ...] = (0, 0, 256, 512, 512, 1024, 2048)
|
|
cutoff_fps_prob: Tuple[float, ...] = (0.1, 0.1, 0.1, 0.1, 0.1, 0.3, 0.2) # sum to 1.0
|
|
|
|
# backbone transformer
|
|
num_enc_layers: int = 0
|
|
hidden_dim: int = 1024
|
|
num_heads: int = 16
|
|
num_dec_layers: int = 24
|
|
dec_hidden_dim: int = 1024
|
|
dec_num_heads: int = 16
|
|
qknorm: bool = True
|
|
qknorm_type: Literal["LayerNorm", "RMSNorm"] = "LayerNorm" # type of qknorm
|
|
salient_attn_mode: Literal["dual_shared", "single", "dual"] = "dual"
|
|
|
|
# query decoder
|
|
fourier_version: Literal["v1", "v2", "v3"] = "v3"
|
|
point_fourier_dim: int = 48 # must be divisible by 6 (sin/cos, x/y/z)
|
|
query_hidden_dim: int = 1024
|
|
query_num_heads: int = 16
|
|
use_flash_query: bool = False
|
|
|
|
# latent code
|
|
latent_size: int = 4096 # == num_fps_point + num_fps_salient_point
|
|
latent_dim: int = 64
|
|
|
|
# loss
|
|
use_ae: bool = False # if true, variance will be ignored, and kl_weight is used as a L2 norm weight
|
|
kl_weight: float = 1e-3
|
|
|
|
# init weights from a pretrained checkpoint
|
|
pretrain_path: Optional[str] = None
|