(args, only_teacher=False, img_size=224)
| 12 | |
| 13 | |
| 14 | def build_model(args, only_teacher=False, img_size=224): |
| 15 | args.arch = args.arch.removesuffix("_memeff") |
| 16 | if "vit" in args.arch: |
| 17 | vit_kwargs = dict( |
| 18 | img_size=img_size, |
| 19 | patch_size=args.patch_size, |
| 20 | init_values=args.layerscale, |
| 21 | ffn_layer=args.ffn_layer, |
| 22 | block_chunks=args.block_chunks, |
| 23 | qkv_bias=args.qkv_bias, |
| 24 | proj_bias=args.proj_bias, |
| 25 | ffn_bias=args.ffn_bias, |
| 26 | num_register_tokens=args.num_register_tokens, |
| 27 | interpolate_offset=args.interpolate_offset, |
| 28 | interpolate_antialias=args.interpolate_antialias, |
| 29 | ) |
| 30 | teacher = vits.__dict__[args.arch](**vit_kwargs) |
| 31 | if only_teacher: |
| 32 | return teacher, teacher.embed_dim |
| 33 | student = vits.__dict__[args.arch]( |
| 34 | **vit_kwargs, |
| 35 | drop_path_rate=args.drop_path_rate, |
| 36 | drop_path_uniform=args.drop_path_uniform, |
| 37 | ) |
| 38 | embed_dim = student.embed_dim |
| 39 | return student, teacher, embed_dim |
| 40 | |
| 41 | |
| 42 | def build_model_from_cfg(cfg, only_teacher=False): |
no outgoing calls
no test coverage detected