|
|
import torch |
|
|
import torch.nn as nn |
|
|
import re |
|
|
import math |
|
|
BatchNorm2d=nn.BatchNorm2d |
|
|
def conv3x3(in_planes, out_planes, stride=1): |
|
|
"""3x3 convolution with padding""" |
|
|
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, |
|
|
padding=1, bias=False) |
|
|
|
|
|
class BasicBlockLayerNorm(nn.Module): |
|
|
expansion = 1 |
|
|
|
|
|
def __init__(self, inplanes, planes,norm_shape, stride=1, downsample=None, dcn=None): |
|
|
super(BasicBlockLayerNorm, self).__init__() |
|
|
self.with_dcn = dcn is not None |
|
|
self.conv1 = conv3x3(inplanes, planes, stride) |
|
|
self.bn1 = nn.LayerNorm(norm_shape) |
|
|
self.relu = nn.ReLU(inplace=True) |
|
|
self.with_modulated_dcn = False |
|
|
if self.with_dcn: |
|
|
fallback_on_stride = dcn.get('fallback_on_stride', False) |
|
|
self.with_modulated_dcn = dcn.get('modulated', False) |
|
|
self.conv2 = conv3x3(planes, planes) |
|
|
if not self.with_dcn or fallback_on_stride: |
|
|
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, |
|
|
padding=1, bias=False) |
|
|
else: |
|
|
raise NotImplementedError |
|
|
self.bn2 = nn.LayerNorm(norm_shape) |
|
|
self.downsample = downsample |
|
|
self.stride = stride |
|
|
|
|
|
def forward(self, x): |
|
|
residual = x |
|
|
|
|
|
out = self.conv1(x) |
|
|
out = self.bn1(out) |
|
|
out = self.relu(out) |
|
|
|
|
|
out = self.conv2(out) |
|
|
if not self.with_dcn: |
|
|
out = self.conv2(out) |
|
|
elif self.with_modulated_dcn: |
|
|
offset_mask = self.conv2_offset(out) |
|
|
offset = offset_mask[:, :18, :, :] |
|
|
mask = offset_mask[:, -9:, :, :].sigmoid() |
|
|
out = self.conv2(out, offset, mask) |
|
|
else: |
|
|
offset = self.conv2_offset(out) |
|
|
out = self.conv2(out, offset) |
|
|
out = self.bn2(out) |
|
|
|
|
|
if self.downsample is not None: |
|
|
residual = self.downsample(x) |
|
|
|
|
|
out += residual |
|
|
out = self.relu(out) |
|
|
|
|
|
return out |
|
|
|
|
|
class BasicBlock(nn.Module): |
|
|
expansion = 1 |
|
|
|
|
|
def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None): |
|
|
super(BasicBlock, self).__init__() |
|
|
self.with_dcn = dcn is not None |
|
|
self.conv1 = conv3x3(inplanes, planes, stride) |
|
|
self.bn1 = nn.BatchNorm2d(planes) |
|
|
self.relu = nn.ReLU(inplace=True) |
|
|
self.with_modulated_dcn = False |
|
|
if self.with_dcn: |
|
|
fallback_on_stride = dcn.get('fallback_on_stride', False) |
|
|
self.with_modulated_dcn = dcn.get('modulated', False) |
|
|
self.conv2 = conv3x3(planes, planes) |
|
|
if not self.with_dcn or fallback_on_stride: |
|
|
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, |
|
|
padding=1, bias=False) |
|
|
else: |
|
|
raise NotImplementedError |
|
|
self.bn2 = nn.BatchNorm2d(planes) |
|
|
self.downsample = downsample |
|
|
self.stride = stride |
|
|
|
|
|
def forward(self, x): |
|
|
residual = x |
|
|
|
|
|
out = self.conv1(x) |
|
|
out = self.bn1(out) |
|
|
out = self.relu(out) |
|
|
|
|
|
out = self.conv2(out) |
|
|
if not self.with_dcn: |
|
|
out = self.conv2(out) |
|
|
elif self.with_modulated_dcn: |
|
|
offset_mask = self.conv2_offset(out) |
|
|
offset = offset_mask[:, :18, :, :] |
|
|
mask = offset_mask[:, -9:, :, :].sigmoid() |
|
|
out = self.conv2(out, offset, mask) |
|
|
else: |
|
|
offset = self.conv2_offset(out) |
|
|
out = self.conv2(out, offset) |
|
|
out = self.bn2(out) |
|
|
|
|
|
if self.downsample is not None: |
|
|
residual = self.downsample(x) |
|
|
|
|
|
out += residual |
|
|
out = self.relu(out) |
|
|
|
|
|
return out |
|
|
|
|
|
class BasicBlockWOnorm(nn.Module): |
|
|
expansion = 1 |
|
|
|
|
|
def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None): |
|
|
super(BasicBlockWOnorm, self).__init__() |
|
|
self.with_dcn = dcn is not None |
|
|
self.conv1 = conv3x3(inplanes, planes, stride) |
|
|
self.bn1 = BatchNorm2d(planes) |
|
|
self.relu = nn.ReLU(inplace=True) |
|
|
self.with_modulated_dcn = False |
|
|
if self.with_dcn: |
|
|
fallback_on_stride = dcn.get('fallback_on_stride', False) |
|
|
self.with_modulated_dcn = dcn.get('modulated', False) |
|
|
self.conv2 = conv3x3(planes, planes) |
|
|
if not self.with_dcn or fallback_on_stride: |
|
|
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, |
|
|
padding=1, bias=False) |
|
|
else: |
|
|
raise NotImplementedError |
|
|
self.bn2 = BatchNorm2d(planes) |
|
|
self.downsample = downsample |
|
|
self.stride = stride |
|
|
|
|
|
def forward(self, x): |
|
|
residual = x |
|
|
|
|
|
out = self.conv1(x) |
|
|
out = self.bn1(out) |
|
|
out = self.relu(out) |
|
|
|
|
|
out = self.conv2(out) |
|
|
if not self.with_dcn: |
|
|
out = self.conv2(out) |
|
|
elif self.with_modulated_dcn: |
|
|
offset_mask = self.conv2_offset(out) |
|
|
offset = offset_mask[:, :18, :, :] |
|
|
mask = offset_mask[:, -9:, :, :].sigmoid() |
|
|
out = self.conv2(out, offset, mask) |
|
|
else: |
|
|
offset = self.conv2_offset(out) |
|
|
out = self.conv2(out, offset) |
|
|
out = self.bn2(out) |
|
|
|
|
|
if self.downsample is not None: |
|
|
residual = self.downsample(x) |
|
|
|
|
|
out += residual |
|
|
out = self.relu(out) |
|
|
|
|
|
return out |
|
|
|
|
|
class ResNetWOnorm(nn.Module): |
|
|
def __init__(self, |
|
|
dcn=None, |
|
|
out_dim=4096): |
|
|
print('using resnet without batchnorm') |
|
|
self.dcn = dcn |
|
|
self.inplanes = 256 |
|
|
super(ResNetWOnorm, self).__init__() |
|
|
self.layer1 = self._make_layer( |
|
|
BasicBlockWOnorm, 1024, 1, stride=2, dcn=dcn) |
|
|
self.layer2 = self._make_layer( |
|
|
BasicBlockWOnorm, 4096, 1, stride=2, dcn=dcn) |
|
|
self.fc = nn.Linear(4096, out_dim) |
|
|
|
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
|
|
m.weight.data.normal_(0, math.sqrt(2. / n)) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
m.weight.data.fill_(1) |
|
|
m.bias.data.zero_() |
|
|
|
|
|
|
|
|
def _make_layer(self, block, planes, blocks, stride=1, dcn=None): |
|
|
downsample = None |
|
|
if stride != 1 or self.inplanes != planes * block.expansion: |
|
|
downsample = nn.Sequential( |
|
|
nn.Conv2d(self.inplanes, planes * block.expansion, |
|
|
kernel_size=1, stride=stride, bias=False), |
|
|
BatchNorm2d(planes * block.expansion), |
|
|
) |
|
|
|
|
|
layers = [] |
|
|
layers.append(block(self.inplanes, planes, |
|
|
stride, downsample, dcn=dcn)) |
|
|
self.inplanes = planes * block.expansion |
|
|
for i in range(1, blocks): |
|
|
layers.append(block(self.inplanes, planes, dcn=dcn)) |
|
|
|
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
x = self.layer1(x) |
|
|
x = self.layer2(x) |
|
|
|
|
|
x = x.reshape(x.shape[0],x.shape[1],-1) |
|
|
x = x.permute(0,2,1) |
|
|
|
|
|
|
|
|
x = self.fc(x) |
|
|
|
|
|
return x |
|
|
|
|
|
class ResNetLayerNorm(nn.Module): |
|
|
def __init__(self, |
|
|
dcn=None, |
|
|
out_dim=4096): |
|
|
print('using resnet with layernorm') |
|
|
self.dcn = dcn |
|
|
self.inplanes = 256 |
|
|
h,w = 64,64 |
|
|
super(ResNetLayerNorm, self).__init__() |
|
|
self.layer1 = self._make_layer( |
|
|
BasicBlockLayerNorm, 1024, 1,[1024,32,32], stride=2, dcn=dcn) |
|
|
self.layer2 = self._make_layer( |
|
|
BasicBlockLayerNorm, 4096, 1,[4096,16,16], stride=2, dcn=dcn) |
|
|
self.fc = nn.Linear(4096, out_dim) |
|
|
|
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
|
|
m.weight.data.normal_(0, math.sqrt(2. / n)) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
m.weight.data.fill_(1) |
|
|
m.bias.data.zero_() |
|
|
|
|
|
|
|
|
def _make_layer(self, block, planes, blocks,norm_shape, stride=1, dcn=None): |
|
|
downsample = None |
|
|
if stride != 1 or self.inplanes != planes * block.expansion: |
|
|
downsample = nn.Sequential( |
|
|
nn.Conv2d(self.inplanes, planes * block.expansion, |
|
|
kernel_size=1, stride=stride, bias=False), |
|
|
nn.LayerNorm(norm_shape), |
|
|
) |
|
|
|
|
|
layers = [] |
|
|
layers.append(block(self.inplanes, planes,norm_shape, |
|
|
stride, downsample, dcn=dcn)) |
|
|
self.inplanes = planes * block.expansion |
|
|
for i in range(1, blocks): |
|
|
layers.append(block(self.inplanes, planes, dcn=dcn)) |
|
|
|
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
x = self.layer1(x) |
|
|
x = self.layer2(x) |
|
|
|
|
|
x = x.reshape(x.shape[0],x.shape[1],-1) |
|
|
x = x.permute(0,2,1) |
|
|
|
|
|
|
|
|
x = self.fc(x) |
|
|
|
|
|
return x |
|
|
|
|
|
class ResNet(nn.Module): |
|
|
def __init__(self, |
|
|
dcn=None, |
|
|
out_dim=4096): |
|
|
self.dcn = dcn |
|
|
self.inplanes = 256 |
|
|
super(ResNet, self).__init__() |
|
|
self.layer1 = self._make_layer( |
|
|
BasicBlock, 1024, 1, stride=2, dcn=dcn) |
|
|
self.layer2 = self._make_layer( |
|
|
BasicBlock, 4096, 1, stride=2, dcn=dcn) |
|
|
self.fc = nn.Linear(4096, out_dim) |
|
|
|
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
|
|
m.weight.data.normal_(0, math.sqrt(2. / n)) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
m.weight.data.fill_(1) |
|
|
m.bias.data.zero_() |
|
|
|
|
|
|
|
|
def _make_layer(self, block, planes, blocks, stride=1, dcn=None): |
|
|
downsample = None |
|
|
if stride != 1 or self.inplanes != planes * block.expansion: |
|
|
downsample = nn.Sequential( |
|
|
nn.Conv2d(self.inplanes, planes * block.expansion, |
|
|
kernel_size=1, stride=stride, bias=False), |
|
|
nn.BatchNorm2d(planes * block.expansion), |
|
|
) |
|
|
|
|
|
layers = [] |
|
|
layers.append(block(self.inplanes, planes, |
|
|
stride, downsample, dcn=dcn)) |
|
|
self.inplanes = planes * block.expansion |
|
|
for i in range(1, blocks): |
|
|
layers.append(block(self.inplanes, planes, dcn=dcn)) |
|
|
|
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
x = self.layer1(x) |
|
|
x = self.layer2(x) |
|
|
|
|
|
x = x.reshape(x.shape[0],x.shape[1],-1) |
|
|
x = x.permute(0,2,1) |
|
|
|
|
|
|
|
|
x = self.fc(x) |
|
|
|
|
|
return x |
|
|
|
|
|
class ResNetSwin(nn.Module): |
|
|
def __init__(self, |
|
|
dcn=None, |
|
|
input_dim=1024, |
|
|
out_dim=4096): |
|
|
self.dcn = dcn |
|
|
self.inplanes = input_dim |
|
|
super(ResNetSwin, self).__init__() |
|
|
self.layer1 = self._make_layer( |
|
|
BasicBlock, 2048, 1, stride=2, dcn=dcn) |
|
|
self.fc = nn.Linear(2048, out_dim) |
|
|
|
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
|
|
m.weight.data.normal_(0, math.sqrt(2. / n)) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
m.weight.data.fill_(1) |
|
|
m.bias.data.zero_() |
|
|
|
|
|
|
|
|
def _make_layer(self, block, planes, blocks, stride=1, dcn=None): |
|
|
downsample = None |
|
|
if stride != 1 or self.inplanes != planes * block.expansion: |
|
|
downsample = nn.Sequential( |
|
|
nn.Conv2d(self.inplanes, planes * block.expansion, |
|
|
kernel_size=1, stride=stride, bias=False), |
|
|
nn.BatchNorm2d(planes * block.expansion), |
|
|
) |
|
|
|
|
|
layers = [] |
|
|
layers.append(block(self.inplanes, planes, |
|
|
stride, downsample, dcn=dcn)) |
|
|
self.inplanes = planes * block.expansion |
|
|
for i in range(1, blocks): |
|
|
layers.append(block(self.inplanes, planes, dcn=dcn)) |
|
|
|
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
x = self.layer1(x) |
|
|
|
|
|
x = x.reshape(x.shape[0],x.shape[1],-1) |
|
|
x = x.permute(0,2,1) |
|
|
|
|
|
|
|
|
x = self.fc(x) |
|
|
|
|
|
return x |
|
|
|
|
|
class IdentityMap(nn.Module): |
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
|
|
|
def forward(self, x, *args, **kwargs): |
|
|
return x |
|
|
|
|
|
@property |
|
|
def config(self): |
|
|
return {"mm_projector_type": 'identity'} |
|
|
|
|
|
|
|
|
class SimpleResBlock(nn.Module): |
|
|
def __init__(self, channels): |
|
|
super().__init__() |
|
|
self.pre_norm = nn.LayerNorm(channels) |
|
|
|
|
|
self.proj = nn.Sequential( |
|
|
nn.Linear(channels, channels), |
|
|
nn.GELU(), |
|
|
nn.Linear(channels, channels) |
|
|
) |
|
|
def forward(self, x): |
|
|
x = self.pre_norm(x) |
|
|
return x + self.proj(x) |
|
|
|
|
|
|
|
|
def build_vision_projector(config, delay_load=False, **kwargs): |
|
|
projector_type = getattr(config, 'mm_projector_type', 'linear') |
|
|
|
|
|
print("projector_type:", projector_type) |
|
|
if projector_type == 'linear': |
|
|
return nn.Linear(config.mm_hidden_size, config.hidden_size) |
|
|
|
|
|
if projector_type == 'conv': |
|
|
with_norm = getattr(config, 'with_norm', True) |
|
|
with_layernorm = getattr(config, 'with_layernorm', True) |
|
|
out_dim = getattr(config,'projector_outdim',4096) |
|
|
|
|
|
if with_layernorm: |
|
|
return ResNetLayerNorm(out_dim=out_dim) |
|
|
if with_norm: |
|
|
return ResNet(out_dim=out_dim) |
|
|
else: |
|
|
return ResNetWOnorm(out_dim=out_dim) |
|
|
if projector_type == 'swin_conv': |
|
|
out_dim = getattr(config,'projector_outdim',4096) |
|
|
input_dim = getattr(config,'mm_input_embeds',1024) |
|
|
return ResNetSwin(input_dim=input_dim,out_dim=out_dim) |
|
|
|
|
|
mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type) |
|
|
print("mlp_gelu_match:", mlp_gelu_match) |
|
|
if mlp_gelu_match: |
|
|
mlp_depth = int(mlp_gelu_match.group(1)) |
|
|
modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)] |
|
|
for _ in range(1, mlp_depth): |
|
|
modules.append(nn.GELU()) |
|
|
modules.append(nn.Linear(config.hidden_size, config.hidden_size)) |
|
|
return nn.Sequential(*modules) |
|
|
|
|
|
if projector_type == 'identity': |
|
|
return IdentityMap() |
|
|
|
|
|
raise ValueError(f'Unknown projector type: {projector_type}') |
|
|
|
|
|
if __name__ == '__main__': |
|
|
class Config: |
|
|
def __init__(self): |
|
|
self.mm_projector_type = 'conv' |
|
|
self.with_layernorm = True |
|
|
self.with_norm = False |
|
|
config = Config() |
|
|
net = build_vision_projector(config) |
|
|
image = torch.randn((4,256,64,64)) |
|
|
print(net) |
|
|
print(net(image).shape) |
