You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
213 lines
9.0 KiB
213 lines
9.0 KiB
2 years ago
|
# Copyright (c) ByteDance, Inc. and its 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.
|
||
|
|
#
|
||
|
|
# This file is basically a copy to: https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
from timm.models.layers import trunc_normal_
|
||
|
|
from timm.models.registry import register_model
|
||
|
|
|
||
|
|
from encoder import SparseConvNeXtBlock, SparseConvNeXtLayerNorm
|
||
|
|
|
||
|
|
|
||
|
|
class ConvNeXt(nn.Module):
|
||
|
|
r""" ConvNeXt
|
||
|
|
A PyTorch impl of : `A ConvNet for the 2020s` -
|
||
|
|
https://arxiv.org/pdf/2201.03545.pdf
|
||
|
|
Args:
|
||
|
|
in_chans (int): Number of input image channels. Default: 3
|
||
|
|
num_classes (int): Number of classes for classification head. Default: 1000
|
||
|
|
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||
|
|
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||
|
|
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||
|
|
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||
|
|
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, in_chans=3, num_classes=1000,
|
||
|
|
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0.,
|
||
|
|
layer_scale_init_value=1e-6, head_init_scale=1., global_pool='avg',
|
||
|
|
sparse=True,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
|
||
|
|
stem = nn.Sequential(
|
||
|
|
nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
|
||
|
|
SparseConvNeXtLayerNorm(dims[0], eps=1e-6, data_format="channels_first", sparse=sparse)
|
||
|
|
)
|
||
|
|
self.downsample_layers.append(stem)
|
||
|
|
for i in range(3):
|
||
|
|
downsample_layer = nn.Sequential(
|
||
|
|
SparseConvNeXtLayerNorm(dims[i], eps=1e-6, data_format="channels_first", sparse=sparse),
|
||
|
|
nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
|
||
|
|
)
|
||
|
|
self.downsample_layers.append(downsample_layer)
|
||
|
|
|
||
|
|
self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
|
||
|
|
self.drop_path_rate = drop_path_rate
|
||
|
|
self.layer_scale_init_value = layer_scale_init_value
|
||
|
|
dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
|
||
|
|
cur = 0
|
||
|
|
for i in range(4):
|
||
|
|
stage = nn.Sequential(
|
||
|
|
*[SparseConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j],
|
||
|
|
layer_scale_init_value=layer_scale_init_value, sparse=sparse) for j in range(depths[i])]
|
||
|
|
)
|
||
|
|
self.stages.append(stage)
|
||
|
|
cur += depths[i]
|
||
|
|
self.depths = depths
|
||
|
|
|
||
|
|
self.apply(self._init_weights)
|
||
|
|
if num_classes > 0:
|
||
|
|
self.norm = SparseConvNeXtLayerNorm(dims[-1], eps=1e-6, sparse=False) # final norm layer for LE/FT; should not be sparse
|
||
|
|
self.fc = nn.Linear(dims[-1], num_classes)
|
||
|
|
# self.fc.weight.data.mul_(head_init_scale) # todo: perform this outside
|
||
|
|
# self.fc.bias.data.mul_(head_init_scale) # todo: perform this outside
|
||
|
|
else:
|
||
|
|
self.norm = nn.Identity()
|
||
|
|
self.fc = nn.Identity()
|
||
|
|
|
||
|
|
self.with_pooling = len(global_pool) > 0
|
||
|
|
|
||
|
|
def _init_weights(self, m):
|
||
|
|
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||
|
|
trunc_normal_(m.weight, std=.02)
|
||
|
|
nn.init.constant_(m.bias, 0)
|
||
|
|
|
||
|
|
def forward_features(self, x, pyramid: int): # pyramid: 0, 1, 2, 3, 4
|
||
|
|
ls = []
|
||
|
|
for i in range(4):
|
||
|
|
x = self.downsample_layers[i](x)
|
||
|
|
x = self.stages[i](x)
|
||
|
|
if pyramid:
|
||
|
|
ls.append(x)
|
||
|
|
|
||
|
|
if pyramid:
|
||
|
|
for i in range(len(ls)-pyramid-1, -1, -1):
|
||
|
|
del ls[i]
|
||
|
|
return [None] * (4 - pyramid) + ls
|
||
|
|
else:
|
||
|
|
if self.with_pooling:
|
||
|
|
x = x.mean([-2, -1]) # global average pooling, (N, C, H, W) -> (N, C)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def forward(self, x, pyramid=0):
|
||
|
|
if pyramid == 0:
|
||
|
|
x = self.forward_features(x, pyramid=pyramid)
|
||
|
|
x = self.fc(self.norm(x))
|
||
|
|
return x
|
||
|
|
else:
|
||
|
|
return self.forward_features(x, pyramid=pyramid)
|
||
|
|
|
||
|
|
def get_classifier(self):
|
||
|
|
return self.fc
|
||
|
|
|
||
|
|
def extra_repr(self):
|
||
|
|
return f'drop_path_rate={self.drop_path_rate}, layer_scale_init_value={self.layer_scale_init_value:g}'
|
||
|
|
|
||
|
|
def get_layer_id_and_scale_exp(self, para_name: str):
|
||
|
|
N = 12 if self.depths[-2] > 9 else 6
|
||
|
|
if para_name.startswith("downsample_layers"):
|
||
|
|
stage_id = int(para_name.split('.')[1])
|
||
|
|
if stage_id == 0:
|
||
|
|
layer_id = 0
|
||
|
|
elif stage_id == 1 or stage_id == 2:
|
||
|
|
layer_id = stage_id + 1
|
||
|
|
else: # stage_id == 3:
|
||
|
|
layer_id = N
|
||
|
|
elif para_name.startswith("stages"):
|
||
|
|
stage_id = int(para_name.split('.')[1])
|
||
|
|
block_id = int(para_name.split('.')[2])
|
||
|
|
if stage_id == 0 or stage_id == 1:
|
||
|
|
layer_id = stage_id + 1
|
||
|
|
elif stage_id == 2:
|
||
|
|
layer_id = 3 + block_id // 3
|
||
|
|
else: # stage_id == 3:
|
||
|
|
layer_id = N
|
||
|
|
else:
|
||
|
|
layer_id = N + 1 # after backbone
|
||
|
|
|
||
|
|
return layer_id, N + 1 - layer_id
|
||
|
|
|
||
|
|
|
||
|
|
model_urls = {
|
||
|
|
"convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
|
||
|
|
"convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
|
||
|
|
"convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
|
||
|
|
"convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
|
||
|
|
"convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
|
||
|
|
"convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
|
||
|
|
"convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
|
||
|
|
"convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
|
||
|
|
"convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@register_model
|
||
|
|
def convnext_tiny(pretrained=False, in_22k=False, **kwargs):
|
||
|
|
model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
|
||
|
|
if pretrained:
|
||
|
|
url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
|
||
|
|
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
|
||
|
|
model.load_state_dict(checkpoint["model"])
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
@register_model
|
||
|
|
def convnext_small(pretrained=False, in_22k=False, **kwargs):
|
||
|
|
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
|
||
|
|
if pretrained:
|
||
|
|
url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
|
||
|
|
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||
|
|
model.load_state_dict(checkpoint["model"])
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
@register_model
|
||
|
|
def convnext_base(pretrained=False, in_22k=False, **kwargs):
|
||
|
|
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
|
||
|
|
if pretrained:
|
||
|
|
url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
|
||
|
|
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||
|
|
model.load_state_dict(checkpoint["model"])
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
@register_model
|
||
|
|
def convnext_large(pretrained=False, in_22k=False, **kwargs):
|
||
|
|
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
|
||
|
|
if pretrained:
|
||
|
|
url = model_urls['convnext_large_22k'] if in_22k else model_urls['convnext_large_1k']
|
||
|
|
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||
|
|
model.load_state_dict(checkpoint["model"])
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
@register_model
|
||
|
|
def convnext_xlarge(pretrained=False, in_22k=False, **kwargs):
|
||
|
|
model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
|
||
|
|
if pretrained:
|
||
|
|
assert in_22k, "only ImageNet-22K pre-trained ConvNeXt-XL is available; please set in_22k=True"
|
||
|
|
url = model_urls['convnext_xlarge_22k']
|
||
|
|
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
|
||
|
|
model.load_state_dict(checkpoint["model"])
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
from timm.models import create_model
|
||
|
|
|
||
|
|
c = create_model('convnext_small', sparse=False)
|
||
|
|
with torch.no_grad():
|
||
|
|
x = torch.rand(2, 3, 224, 224)
|
||
|
|
print(c(x).shape)
|
||
|
|
print([None if f is None else f.shape for f in c(x, pyramid=1)])
|
||
|
|
print([None if f is None else f.shape for f in c(x, pyramid=2)])
|
||
|
|
print([None if f is None else f.shape for f in c(x, pyramid=3)])
|
||
|
|
print([None if f is None else f.shape for f in c(x, pyramid=4)])
|