Spaces:
Running
on
Zero
Running
on
Zero
SunderAli17
commited on
Commit
•
47de9d9
1
Parent(s):
111e24f
Create modified_resnet.py
Browse files- eva_clip/modified_resnet.py +181 -0
eva_clip/modified_resnet.py
ADDED
@@ -0,0 +1,181 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from collections import OrderedDict
|
2 |
+
|
3 |
+
import torch
|
4 |
+
from torch import nn
|
5 |
+
from torch.nn import functional as F
|
6 |
+
|
7 |
+
from eva_clip.utils import freeze_batch_norm_2d
|
8 |
+
|
9 |
+
|
10 |
+
class Bottleneck(nn.Module):
|
11 |
+
expansion = 4
|
12 |
+
|
13 |
+
def __init__(self, inplanes, planes, stride=1):
|
14 |
+
super().__init__()
|
15 |
+
|
16 |
+
# all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
|
17 |
+
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
|
18 |
+
self.bn1 = nn.BatchNorm2d(planes)
|
19 |
+
self.act1 = nn.ReLU(inplace=True)
|
20 |
+
|
21 |
+
self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
|
22 |
+
self.bn2 = nn.BatchNorm2d(planes)
|
23 |
+
self.act2 = nn.ReLU(inplace=True)
|
24 |
+
|
25 |
+
self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
|
26 |
+
|
27 |
+
self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
|
28 |
+
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
29 |
+
self.act3 = nn.ReLU(inplace=True)
|
30 |
+
|
31 |
+
self.downsample = None
|
32 |
+
self.stride = stride
|
33 |
+
|
34 |
+
if stride > 1 or inplanes != planes * Bottleneck.expansion:
|
35 |
+
# downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
|
36 |
+
self.downsample = nn.Sequential(OrderedDict([
|
37 |
+
("-1", nn.AvgPool2d(stride)),
|
38 |
+
("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
|
39 |
+
("1", nn.BatchNorm2d(planes * self.expansion))
|
40 |
+
]))
|
41 |
+
|
42 |
+
def forward(self, x: torch.Tensor):
|
43 |
+
identity = x
|
44 |
+
|
45 |
+
out = self.act1(self.bn1(self.conv1(x)))
|
46 |
+
out = self.act2(self.bn2(self.conv2(out)))
|
47 |
+
out = self.avgpool(out)
|
48 |
+
out = self.bn3(self.conv3(out))
|
49 |
+
|
50 |
+
if self.downsample is not None:
|
51 |
+
identity = self.downsample(x)
|
52 |
+
|
53 |
+
out += identity
|
54 |
+
out = self.act3(out)
|
55 |
+
return out
|
56 |
+
|
57 |
+
|
58 |
+
class AttentionPool2d(nn.Module):
|
59 |
+
def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
|
60 |
+
super().__init__()
|
61 |
+
self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
|
62 |
+
self.k_proj = nn.Linear(embed_dim, embed_dim)
|
63 |
+
self.q_proj = nn.Linear(embed_dim, embed_dim)
|
64 |
+
self.v_proj = nn.Linear(embed_dim, embed_dim)
|
65 |
+
self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
|
66 |
+
self.num_heads = num_heads
|
67 |
+
|
68 |
+
def forward(self, x):
|
69 |
+
x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC
|
70 |
+
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC
|
71 |
+
x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
|
72 |
+
x, _ = F.multi_head_attention_forward(
|
73 |
+
query=x, key=x, value=x,
|
74 |
+
embed_dim_to_check=x.shape[-1],
|
75 |
+
num_heads=self.num_heads,
|
76 |
+
q_proj_weight=self.q_proj.weight,
|
77 |
+
k_proj_weight=self.k_proj.weight,
|
78 |
+
v_proj_weight=self.v_proj.weight,
|
79 |
+
in_proj_weight=None,
|
80 |
+
in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
|
81 |
+
bias_k=None,
|
82 |
+
bias_v=None,
|
83 |
+
add_zero_attn=False,
|
84 |
+
dropout_p=0.,
|
85 |
+
out_proj_weight=self.c_proj.weight,
|
86 |
+
out_proj_bias=self.c_proj.bias,
|
87 |
+
use_separate_proj_weight=True,
|
88 |
+
training=self.training,
|
89 |
+
need_weights=False
|
90 |
+
)
|
91 |
+
|
92 |
+
return x[0]
|
93 |
+
|
94 |
+
|
95 |
+
class ModifiedResNet(nn.Module):
|
96 |
+
"""
|
97 |
+
A ResNet class that is similar to torchvision's but contains the following changes:
|
98 |
+
- There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
|
99 |
+
- Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
|
100 |
+
- The final pooling layer is a QKV attention instead of an average pool
|
101 |
+
"""
|
102 |
+
|
103 |
+
def __init__(self, layers, output_dim, heads, image_size=224, width=64):
|
104 |
+
super().__init__()
|
105 |
+
self.output_dim = output_dim
|
106 |
+
self.image_size = image_size
|
107 |
+
|
108 |
+
# the 3-layer stem
|
109 |
+
self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
|
110 |
+
self.bn1 = nn.BatchNorm2d(width // 2)
|
111 |
+
self.act1 = nn.ReLU(inplace=True)
|
112 |
+
self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
|
113 |
+
self.bn2 = nn.BatchNorm2d(width // 2)
|
114 |
+
self.act2 = nn.ReLU(inplace=True)
|
115 |
+
self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
|
116 |
+
self.bn3 = nn.BatchNorm2d(width)
|
117 |
+
self.act3 = nn.ReLU(inplace=True)
|
118 |
+
self.avgpool = nn.AvgPool2d(2)
|
119 |
+
|
120 |
+
# residual layers
|
121 |
+
self._inplanes = width # this is a *mutable* variable used during construction
|
122 |
+
self.layer1 = self._make_layer(width, layers[0])
|
123 |
+
self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
|
124 |
+
self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
|
125 |
+
self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
|
126 |
+
|
127 |
+
embed_dim = width * 32 # the ResNet feature dimension
|
128 |
+
self.attnpool = AttentionPool2d(image_size // 32, embed_dim, heads, output_dim)
|
129 |
+
|
130 |
+
self.init_parameters()
|
131 |
+
|
132 |
+
def _make_layer(self, planes, blocks, stride=1):
|
133 |
+
layers = [Bottleneck(self._inplanes, planes, stride)]
|
134 |
+
|
135 |
+
self._inplanes = planes * Bottleneck.expansion
|
136 |
+
for _ in range(1, blocks):
|
137 |
+
layers.append(Bottleneck(self._inplanes, planes))
|
138 |
+
|
139 |
+
return nn.Sequential(*layers)
|
140 |
+
|
141 |
+
def init_parameters(self):
|
142 |
+
if self.attnpool is not None:
|
143 |
+
std = self.attnpool.c_proj.in_features ** -0.5
|
144 |
+
nn.init.normal_(self.attnpool.q_proj.weight, std=std)
|
145 |
+
nn.init.normal_(self.attnpool.k_proj.weight, std=std)
|
146 |
+
nn.init.normal_(self.attnpool.v_proj.weight, std=std)
|
147 |
+
nn.init.normal_(self.attnpool.c_proj.weight, std=std)
|
148 |
+
|
149 |
+
for resnet_block in [self.layer1, self.layer2, self.layer3, self.layer4]:
|
150 |
+
for name, param in resnet_block.named_parameters():
|
151 |
+
if name.endswith("bn3.weight"):
|
152 |
+
nn.init.zeros_(param)
|
153 |
+
|
154 |
+
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
155 |
+
assert unlocked_groups == 0, 'partial locking not currently supported for this model'
|
156 |
+
for param in self.parameters():
|
157 |
+
param.requires_grad = False
|
158 |
+
if freeze_bn_stats:
|
159 |
+
freeze_batch_norm_2d(self)
|
160 |
+
|
161 |
+
@torch.jit.ignore
|
162 |
+
def set_grad_checkpointing(self, enable=True):
|
163 |
+
# FIXME support for non-transformer
|
164 |
+
pass
|
165 |
+
|
166 |
+
def stem(self, x):
|
167 |
+
x = self.act1(self.bn1(self.conv1(x)))
|
168 |
+
x = self.act2(self.bn2(self.conv2(x)))
|
169 |
+
x = self.act3(self.bn3(self.conv3(x)))
|
170 |
+
x = self.avgpool(x)
|
171 |
+
return x
|
172 |
+
|
173 |
+
def forward(self, x):
|
174 |
+
x = self.stem(x)
|
175 |
+
x = self.layer1(x)
|
176 |
+
x = self.layer2(x)
|
177 |
+
x = self.layer3(x)
|
178 |
+
x = self.layer4(x)
|
179 |
+
x = self.attnpool(x)
|
180 |
+
|
181 |
+
return x
|