(pretrained=None, **kwargs)
| 416 | |
| 417 | |
| 418 | def small(pretrained=None, **kwargs): |
| 419 | model = VisionTransformer( |
| 420 | patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True, |
| 421 | norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs) |
| 422 | if pretrained: |
| 423 | # checkpoint = torch.load('deit_small_patch16_224-cd65a155.pth', map_location="cpu") |
| 424 | # checkpoint = torch.hub.load_state_dict_from_url( |
| 425 | # url="https://dl.fbaipublicfiles.com/deit/deit_small_patch16_224-cd65a155.pth", |
| 426 | # map_location="cpu", check_hash=True |
| 427 | # ) |
| 428 | checkpoint = torch.load(pretrained, map_location="cpu") |
| 429 | model.load_state_dict(checkpoint["model"], strict=False) |
| 430 | return model, 384 |
| 431 | |
| 432 | |
| 433 | def small_dWr(pretrained=None, **kwargs): |
no test coverage detected