This repository has been archived by the owner on Oct 31, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 24
/
hubconf.py
90 lines (71 loc) · 2.86 KB
/
hubconf.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import torch
import resnet
import convnext
dependencies = ["torch", "torchvision"]
def resnet50_alpha0p9(pretrained=True, **kwargs):
model, _ = resnet.resnet50(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/resnet50_alpha0.9.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model
def resnet50_alpha0p75(pretrained=True, **kwargs):
model, _ = resnet.resnet50(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/resnet50_alpha0.75.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model
def convnext_small_alpha0p9(pretrained=True, **kwargs):
model, _ = convnext.convnext_small(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/convnext_small_alpha0.9.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model
def convnext_small_alpha0p75(pretrained=True, **kwargs):
model, _ = convnext.convnext_small(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/convnext_small_alpha0.75.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model
def convnext_base_alpha0p9(pretrained=True, **kwargs):
model, _ = convnext.convnext_base(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/convnext_base_alpha0.9.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model
def convnext_base_alpha0p75(pretrained=True, **kwargs):
model, _ = convnext.convnext_base(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/convnext_base_alpha0.75.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model
def convnext_xlarge_alpha0p75(pretrained=True, **kwargs):
model, _ = convnext.convnext_xlarge(**kwargs)
if pretrained:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/vicregl/convnext_xlarge_alpha0.75.pth",
map_location="cpu",
)
model.load_state_dict(state_dict, strict=True)
return model