# Generated by configen, do not edit.
# See https://github.com/facebookresearch/hydra/tree/master/tools/configen
# fmt: off
# isort:skip_file
# flake8: noqa

from dataclasses import dataclass, field
from conduit.architectures.vit import VitArch
from conduit.models.self_supervised.utils import ResNetArch
from typing import Any
from typing import Optional
from typing import Tuple


@dataclass
class MoCoV2Conf:
    _target_: str = "conduit.models.self_supervised.MoCoV2"
    backbone: Any = ResNetArch.resnet18  # Union[Module, ResNetArch, str]
    head: Any = None  # Optional[Module]
    out_dim: int = 128
    num_negatives: int = 65536
    momentum_teacher: float = 0.999
    temp: float = 0.07
    lr: float = 0.03
    momentum_sgd: float = 0.9
    weight_decay: float = 0.0001
    use_mlp: bool = False
    instance_transforms: Any = None  # Optional[MultiCropTransform]
    batch_transforms: Any = None  # Optional[BatchTransform]
    global_crop_size: Any = None  # Union[Tuple[int, int], int, NoneType]
    local_crop_size: Any = 0.43  # Union[Tuple[float, float], float]
    global_crops_scale: Tuple[float, float] = (0.4, 1.0)
    local_crops_scale: Tuple[float, float] = (0.05, 0.4)
    local_crops_number: int = 0
    eval_epochs: int = 100
    eval_batch_size: Optional[int] = None


@dataclass
class DINOConf:
    _target_: str = "conduit.models.self_supervised.DINO"
    backbone: Any = VitArch.small  # Union[Module, VitArch, VisionTransformer, ResNetArch, str]
    out_dim: int = 65536
    lr: float = 0.0005
    warmup_iters: int = 10
    weight_decay: float = 0.04
    min_lr: float = 1e-06
    weight_decay_final: float = 0.4
    freeze_last_layer: int = 1
    patch_size: int = 16
    drop_path_rate: float = 0.1
    norm_last_layer: bool = True
    use_bn_in_head: bool = False
    momentum_teacher: float = 0.996
    momentum_center: float = 0.9
    teacher_temp: float = 0.04
    warmup_teacher_temp: float = 0.04
    student_temp: float = 0.1
    warmup_teacher_temp_iters: int = 30
    num_eval_blocks: int = 1
    lr_eval: float = 0.0001
    global_crop_size: Any = None  # Union[Tuple[int, int], int, NoneType]
    local_crop_size: Any = 0.43  # Union[Tuple[float, float], float]
    global_crops_scale: Tuple[float, float] = (0.4, 1.0)
    local_crops_scale: Tuple[float, float] = (0.05, 0.4)
    local_crops_number: int = 8
    batch_transforms: Any = None  # Optional[BatchTransform]
    eval_epochs: int = 100
    eval_batch_size: Optional[int] = None
