44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the Apache License, Version 2.0
|
|
# found in the LICENSE file in the root directory of this source tree.
|
|
|
|
import logging
|
|
|
|
from . import vision_transformer as vits
|
|
|
|
|
|
logger = logging.getLogger("dinov2")
|
|
|
|
|
|
def build_model(args, only_teacher=False, img_size=224):
|
|
args.arch = args.arch.removesuffix("_memeff")
|
|
if "vit" in args.arch:
|
|
vit_kwargs = dict(
|
|
img_size=img_size,
|
|
patch_size=args.patch_size,
|
|
init_values=args.layerscale,
|
|
ffn_layer=args.ffn_layer,
|
|
block_chunks=args.block_chunks,
|
|
qkv_bias=args.qkv_bias,
|
|
proj_bias=args.proj_bias,
|
|
ffn_bias=args.ffn_bias,
|
|
num_register_tokens=args.num_register_tokens,
|
|
interpolate_offset=args.interpolate_offset,
|
|
interpolate_antialias=args.interpolate_antialias,
|
|
)
|
|
teacher = vits.__dict__[args.arch](**vit_kwargs)
|
|
if only_teacher:
|
|
return teacher, teacher.embed_dim
|
|
student = vits.__dict__[args.arch](
|
|
**vit_kwargs,
|
|
drop_path_rate=args.drop_path_rate,
|
|
drop_path_uniform=args.drop_path_uniform,
|
|
)
|
|
embed_dim = student.embed_dim
|
|
return student, teacher, embed_dim
|
|
|
|
|
|
def build_model_from_cfg(cfg, only_teacher=False):
|
|
return build_model(cfg.student, only_teacher=only_teacher, img_size=cfg.crops.global_crops_size)
|