| 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) |
