00:00:00
1 介绍
记录写好的预训练的模型的代码
1 MetaFormer
论文名称：MetaFormer Is Actually What You Need for Vision
论文链接：[2111.11418] MetaFormer Is Actually What You Need for Vision
论文代码：sail-sg/poolformer: PoolFormer: MetaFormer Is Actually What You Need for Vision (CVPR 2022 Oral)
参考博客：CVPR 2022 Oral | MetaFormer：证明Transformer的威力源自其整体架构！颜水成团队工作！...-CSDN博客
1.1 代码实现
1.1.1 网络模块代码
1、导包
python
"""
PoolFormer implementation
"""
import os
import copy
import torch
import torch.nn as nn
from torch.nn.modules.batchnorm import _BatchNorm
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.layers import DropPath, trunc_normal_
from timm.layers.helpers import to_2tuple
from typing import Sequence
from functools import partial, reduce
import torchvision
import time
from torch.utils.data import DataLoader
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path
from tqdm import tqdm
# 预训练模型代码链接
def _cfg(url='', **kwargs):
return {
'url': url,
'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
'crop_pct': .95, 'interpolation': 'bicubic',
'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
'classifier': 'head',
**kwargs
}
default_cfgs = {
'poolformer_s': _cfg(crop_pct=0.9),
'poolformer_m': _cfg(crop_pct=0.95),
}
model_urls_for_poolformer = {
"poolformer_s12": "https://github.com/sail-sg/poolformer/releases/download/v1.0/poolformer_s12.pth.tar",
"poolformer_s24": "https://github.com/sail-sg/poolformer/releases/download/v1.0/poolformer_s24.pth.tar",
"poolformer_s36": "https://github.com/sail-sg/poolformer/releases/download/v1.0/poolformer_s36.pth.tar",
"poolformer_m36": "https://github.com/sail-sg/poolformer/releases/download/v1.0/poolformer_m36.pth.tar",
"poolformer_m48": "https://github.com/sail-sg/poolformer/releases/download/v1.0/poolformer_m48.pth.tar",
}2、模块代码
（1）PatchEmbed
python
class PatchEmbed(nn.Module):
"""
Patch Embedding that is implemented by a layer of conv.
Input: tensor in shape [B, C, H, W]
Output: tensor in shape [B, C, H/stride, W/stride]
"""
def __init__(self, patch_size=16, stride=16, padding=0,
in_chans=3, embed_dim=768, norm_layer=None):
super().__init__()
patch_size = to_2tuple(patch_size)
stride = to_2tuple(stride)
padding = to_2tuple(padding)
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size,
stride=stride, padding=padding)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x):
x = self.proj(x)
x = self.norm(x)
return x（2）LayerNormChannel
python
class LayerNormChannel(nn.Module):
"""
LayerNorm only for Channel Dimension.
Input: tensor in shape [B, C, H, W]
"""
def __init__(self, num_channels, eps=1e-05):
super().__init__()
self.weight = nn.Parameter(torch.ones(num_channels))
self.bias = nn.Parameter(torch.zeros(num_channels))
self.eps = eps
def forward(self, x):
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight.unsqueeze(-1).unsqueeze(-1) * x \
+ self.bias.unsqueeze(-1).unsqueeze(-1)
return x（3）GroupNorm
python
class GroupNorm(nn.GroupNorm):
"""
Group Normalization with 1 group.
Input: tensor in shape [B, C, H, W]
"""
def __init__(self, num_channels, **kwargs):
super().__init__(1, num_channels, **kwargs)（4）Pooling
python
class Pooling(nn.Module):
"""
Implementation of pooling for PoolFormer
--pool_size: pooling size
"""
def __init__(self, pool_size=3):
super().__init__()
self.pool = nn.AvgPool2d(
pool_size, stride=1, padding=pool_size//2, count_include_pad=False)
def forward(self, x):
return self.pool(x) - x（5）Mlp
python
class Mlp(nn.Module):
"""
Implementation of MLP with 1*1 convolutions.
Input: tensor with shape [B, C, H, W]
"""
def __init__(self, in_features, hidden_features=None,
out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
self.act = act_layer()
self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
self.drop = nn.Dropout(drop)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Conv2d):
trunc_normal_(m.weight, std=.02)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x（6）PoolFormerBlock
python
class PoolFormerBlock(nn.Module):
"""
Implementation of one PoolFormer block.
--dim: embedding dim
--pool_size: pooling size
--mlp_ratio: mlp expansion ratio
--act_layer: activation
--norm_layer: normalization
--drop: dropout rate
--drop path: Stochastic Depth,
refer to https://arxiv.org/abs/1603.09382
--use_layer_scale, --layer_scale_init_value: LayerScale,
refer to https://arxiv.org/abs/2103.17239
"""
def __init__(self, dim, pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, norm_layer=GroupNorm,
drop=0., drop_path=0.,
use_layer_scale=True, layer_scale_init_value=1e-5):
super().__init__()
self.norm1 = norm_layer(dim)
self.token_mixer = Pooling(pool_size=pool_size)
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
act_layer=act_layer, drop=drop)
# The following two techniques are useful to train deep PoolFormers.
self.drop_path = DropPath(drop_path) if drop_path > 0. \
else nn.Identity()
self.use_layer_scale = use_layer_scale
if use_layer_scale:
self.layer_scale_1 = nn.Parameter(
layer_scale_init_value * torch.ones((dim)), requires_grad=True)
self.layer_scale_2 = nn.Parameter(
layer_scale_init_value * torch.ones((dim)), requires_grad=True)
def forward(self, x):
# print("==============PoolFormer Block=================")
# print("input", x.shape)
if self.use_layer_scale:
# print("==============use_layer_scale======================")
# print("norm1", self.norm1(x).shape)
# print("self.layer_scale_1.unsqueeze(-1).unsqueeze(-1)", self.layer_scale_1.unsqueeze(-1).unsqueeze(-1).shape)
x = x + self.drop_path(
self.layer_scale_1.unsqueeze(-1).unsqueeze(-1)
* self.token_mixer(self.norm1(x)))
# print("token_mixer", x.shape)
x = x + self.drop_path(
self.layer_scale_2.unsqueeze(-1).unsqueeze(-1)
* self.mlp(self.norm2(x)))
# print("mlp", x.shape)
else:
# print("==============not use_layer_scale======================")
x = x + self.drop_path(self.token_mixer(self.norm1(x)))
# print("token_mixer", x.shape)
x = x + self.drop_path(self.mlp(self.norm2(x)))
# print("mlp", x.shape)
return x（7）PoolFormer
python
class PoolFormer(nn.Module):
"""
PoolFormer, the main class of our model
--layers: [x,x,x,x], number of blocks for the 4 stages
--embed_dims, --mlp_ratios, --pool_size: the embedding dims, mlp ratios and
pooling size for the 4 stages
--downsamples: flags to apply downsampling or not
--norm_layer, --act_layer: define the types of normalization and activation
--num_classes: number of classes for the image classification
--in_patch_size, --in_stride, --in_pad: specify the patch embedding
for the input image
--down_patch_size --down_stride --down_pad:
specify the downsample (patch embed.)
--fork_feat: whether output features of the 4 stages, for dense prediction
--init_cfg, --pretrained:
for mmdetection and mmsegmentation to load pretrained weights
"""
def __init__(self, layers, embed_dims=None,
mlp_ratios=None, downsamples=None,
pool_size=3,
norm_layer=GroupNorm, act_layer=nn.GELU,
num_classes=1000,
in_patch_size=7, in_stride=4, in_pad=2,
down_patch_size=3, down_stride=2, down_pad=1,
drop_rate=0., drop_path_rate=0.,
use_layer_scale=True, layer_scale_init_value=1e-5,
fork_feat=False,
init_cfg=None,
pretrained=None,
**kwargs):
super().__init__()
if not fork_feat:
self.num_classes = num_classes
self.fork_feat = fork_feat
self.patch_embed = PatchEmbed(
patch_size=in_patch_size, stride=in_stride, padding=in_pad,
in_chans=3, embed_dim=embed_dims[0])
# set the main block in network
network = []
for i in range(len(layers)):
stage = basic_blocks_for_poolformer(embed_dims[i], i, layers,
pool_size=pool_size, mlp_ratio=mlp_ratios[i],
act_layer=act_layer, norm_layer=norm_layer,
drop_rate=drop_rate,
drop_path_rate=drop_path_rate,
use_layer_scale=use_layer_scale,
layer_scale_init_value=layer_scale_init_value)
network.append(stage)
if i >= len(layers) - 1:
break
if downsamples[i] or embed_dims[i] != embed_dims[i+1]:
# downsampling between two stages
network.append(
PatchEmbed(
patch_size=down_patch_size, stride=down_stride,
padding=down_pad,
in_chans=embed_dims[i], embed_dim=embed_dims[i+1]
)
)
self.network = nn.ModuleList(network)
if self.fork_feat:
# add a norm layer for each output
self.out_indices = [0, 2, 4, 6]
for i_emb, i_layer in enumerate(self.out_indices):
if i_emb == 0 and os.environ.get('FORK_LAST3', None):
# TODO: more elegant way
"""For RetinaNet, `start_level=1`. The first norm layer will not used.
cmd: `FORK_LAST3=1 python -m torch.distributed.launch ...`
"""
layer = nn.Identity()
else:
layer = norm_layer(embed_dims[i_emb])
layer_name = f'norm{i_layer}'
self.add_module(layer_name, layer)
else:
# Classifier head
self.norm = norm_layer(embed_dims[-1])
self.head = nn.Linear(
embed_dims[-1], num_classes) if num_classes > 0 \
else nn.Identity()
self.apply(self.cls_init_weights)
self.init_cfg = copy.deepcopy(init_cfg)
# load pre-trained model
if self.fork_feat and (
self.init_cfg is not None or pretrained is not None):
self.init_weights()
# init for classification
def cls_init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
def get_classifier(self):
return self.head
def reset_classifier(self, num_classes):
self.num_classes = num_classes
self.head = nn.Linear(
self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
def forward_embeddings(self, x):
x = self.patch_embed(x)
return x
def forward_tokens(self, x):
outs = []
# print(f"network:{self.network}")
# print(f"first of tokens:{x.shape}") # first of tokens:torch.Size([8, 64, 28, 28])
# print("==========================================")
for idx, block in enumerate(self.network):
# print(f"idx:{idx}, block:{block}")
x = block(x)
if self.fork_feat and idx in self.out_indices:
print("use the norm")
norm_layer = getattr(self, f'norm{idx}')
x_out = norm_layer(x)
outs.append(x_out)
# print(f"idx : {idx}, x.shape : {x.shape}")
if self.fork_feat:
# output the features of four stages for dense prediction
return outs
# output only the features of last layer for image classification
return x
def forward(self, x):
# print("input", x.shape) # input torch.Size([8, 3, 112, 112])
# input embedding
x = self.forward_embeddings(x)
# print("forward_embeddings", x.shape) # forward_embeddings torch.Size([8, 64, 28, 28])
# through backbone
x = self.forward_tokens(x)
# print("forward_tokens", x.shape) # forward_tokens torch.Size([8, 512, 4, 4])
if self.fork_feat:
# otuput features of four stages for dense prediction
return x
x = self.norm(x)
# print("norm", x.shape) # norm torch.Size([8, 512, 4, 4])
# print("mean[-2, -1]", x.mean([-2, -1]).shape) # mean[-2, -1] torch.Size([8, 512])
cls_out = self.head(x.mean([-2, -1]))
# print("cls_out", cls_out.shape) # cls_out torch.Size([8, 10])
# for image classification
return cls_out（8）定义不同层数的Poolformer
python
def basic_blocks_for_poolformer(dim, index, layers,
pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, norm_layer=GroupNorm,
drop_rate=.0, drop_path_rate=0.,
use_layer_scale=True, layer_scale_init_value=1e-5):
"""
generate PoolFormer blocks for a stage
return: PoolFormer blocks
"""
blocks = []
for block_idx in range(layers[index]):
block_dpr = drop_path_rate * (
block_idx + sum(layers[:index])) / (sum(layers) - 1)
blocks.append(PoolFormerBlock(
dim, pool_size=pool_size, mlp_ratio=mlp_ratio,
act_layer=act_layer, norm_layer=norm_layer,
drop=drop_rate, drop_path=block_dpr,
use_layer_scale=use_layer_scale,
layer_scale_init_value=layer_scale_init_value,
))
blocks = nn.Sequential(*blocks)
return blocks
def poolformer_s12(pretrained=False, **kwargs):
"""
PoolFormer-S12 model, Params: 12M
--layers: [x,x,x,x], numbers of layers for the four stages
--embed_dims, --mlp_ratios:
embedding dims and mlp ratios for the four stages
--downsamples: flags to apply downsampling or not in four blocks
"""
layers = [2, 2, 6, 2]
embed_dims = [64, 128, 320, 512]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
**kwargs)
model.default_cfg = default_cfgs['poolformer_s']
if pretrained:
url = model_urls_for_poolformer['poolformer_s12']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint)
return model
def poolformer_s24(pretrained=False, **kwargs):
"""
PoolFormer-S24 model, Params: 21M
"""
layers = [4, 4, 12, 4]
embed_dims = [64, 128, 320, 512]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
**kwargs)
model.default_cfg = default_cfgs['poolformer_s']
if pretrained:
url = model_urls_for_poolformer['poolformer_s24']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint)
return model
def poolformer_s36(pretrained=False, **kwargs):
"""
PoolFormer-S36 model, Params: 31M
"""
layers = [6, 6, 18, 6]
embed_dims = [64, 128, 320, 512]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
layer_scale_init_value=1e-6,
**kwargs)
model.default_cfg = default_cfgs['poolformer_s']
if pretrained:
url = model_urls_for_poolformer['poolformer_s36']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint)
return model
def poolformer_m36(pretrained=False, **kwargs):
"""
PoolFormer-M36 model, Params: 56M
"""
layers = [6, 6, 18, 6]
embed_dims = [96, 192, 384, 768]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
layer_scale_init_value=1e-6,
**kwargs)
model.default_cfg = default_cfgs['poolformer_m']
if pretrained:
url = model_urls_for_poolformer['poolformer_m36']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint)
return model
def poolformer_m48(pretrained=False, **kwargs):
"""
PoolFormer-M48 model, Params: 73M
"""
layers = [8, 8, 24, 8]
embed_dims = [96, 192, 384, 768]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
layer_scale_init_value=1e-6,
**kwargs)
model.default_cfg = default_cfgs['poolformer_m']
if pretrained:
url = model_urls_for_poolformer['poolformer_m48']
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
model.load_state_dict(checkpoint)
return model1.1.2 测试网络输入输出
python
if __name__ == "__main__":
model = poolformer_s12(num_classes=10)
# model = poolformer_s24(num_classes=10)
# model = poolformer_s36(num_classes=10)
# model = poolformer_m36(num_classes=10)
# model = poolformer_m48(num_classes=10)
print("参数量", sum(p.numel() for p in model.parameters()))
x = torch.randn(8, 3, 112, 112)
print("input", x.shape)
y = model(x)
print("output", y.shape) # torch.Size([8, 10])
"""
参数量 11407306
input torch.Size([8, 3, 112, 112])
output torch.Size([8, 10])
"""1.1.3 PoolFormer 3D版本
python
"""
PoolFormer implementation For 3D
"""
import os
import copy
import torch
import torch.nn as nn
from timm.models.layers import DropPath, trunc_normal_
class PatchEmbed(nn.Module):
"""
Patch Embedding that is implemented by a layer of conv.
Input: tensor in shape [B, C, D, H, W]
Output: tensor in shape [B, C, D/stride, H/stride, W/stride]
"""
def __init__(self, patch_size=16, stride=16, padding=0,
in_chans=3, embed_dim=768, norm_layer=None):
super().__init__()
# patch_size = to_3tuple(patch_size)
patch_size = (patch_size, patch_size, patch_size)
# stride = to_3tuple(stride)
stride = (stride, stride, stride)
# padding = to_3tuple(padding)
padding = (padding, padding, padding)
self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size,
stride=stride, padding=padding)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x):
x = self.proj(x)
x = self.norm(x)
return x
class LayerNormChannel(nn.Module):
"""
LayerNorm only for Channel Dimension.
Input: tensor in shape [B, C, D, H, W]
"""
def __init__(self, num_channels, eps=1e-05):
super().__init__()
self.weight = nn.Parameter(torch.ones(num_channels))
self.bias = nn.Parameter(torch.zeros(num_channels))
self.eps = eps
def forward(self, x):
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight.view(1, -1, 1, 1, 1) * x + self.bias.view(1, -1, 1, 1, 1)
return x
#
class GroupNorm(nn.GroupNorm):
"""
Group Normalization with 1 group.
Input: tensor in shape [B, C, D, H, W]
"""
def __init__(self, num_channels, **kwargs):
super().__init__(1, num_channels, **kwargs)
class Pooling(nn.Module):
"""
Implementation of pooling for PoolFormer
--pool_size: pooling size
"""
def __init__(self, pool_size=3):
super().__init__()
self.pool = nn.AvgPool3d(
pool_size, stride=1, padding=pool_size//2, count_include_pad=False)
def forward(self, x):
return self.pool(x) - x
class Mlp(nn.Module):
"""
Implementation of MLP with 1*1 convolutions.
Input: tensor with shape [B, C, D, H, W]
"""
def __init__(self, in_features, hidden_features=None,
out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Conv3d(in_features, hidden_features, 1)
self.act = act_layer()
self.fc2 = nn.Conv3d(hidden_features, out_features, 1)
self.drop = nn.Dropout(drop)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Conv2d):
trunc_normal_(m.weight, std=.02)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class PoolFormerBlock(nn.Module):
"""
Implementation of one PoolFormer block.
--dim: embedding dim
--pool_size: pooling size
--mlp_ratio: mlp expansion ratio
--act_layer: activation
--norm_layer: normalization
--drop: dropout rate
--drop path: Stochastic Depth,
refer to https://arxiv.org/abs/1603.09382
--use_layer_scale, --layer_scale_init_value: LayerScale,
refer to https://arxiv.org/abs/2103.17239
"""
def __init__(self, dim, pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, norm_layer=GroupNorm,
drop=0., drop_path=0.,
use_layer_scale=True, layer_scale_init_value=1e-5):
super().__init__()
self.norm1 = norm_layer(dim)
self.token_mixer = Pooling(pool_size=pool_size)
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
act_layer=act_layer, drop=drop)
# The following two techniques are useful to train deep PoolFormers.
self.drop_path = DropPath(drop_path) if drop_path > 0. \
else nn.Identity()
self.use_layer_scale = use_layer_scale
if use_layer_scale:
self.layer_scale_1 = nn.Parameter(
layer_scale_init_value * torch.ones((dim)), requires_grad=True)
self.layer_scale_2 = nn.Parameter(
layer_scale_init_value * torch.ones((dim)), requires_grad=True)
def forward(self, x):
# print("==============PoolFormer Block=================")
# print("input", x.shape)
if self.use_layer_scale:
# print("==============use_layer_scale======================")
# print("norm1", self.norm1(x).shape)
# print("self.layer_scale_1.unsqueeze(-1).unsqueeze(-1)", self.layer_scale_1.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).shape)
x = x + self.drop_path(
self.layer_scale_1.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
* self.token_mixer(self.norm1(x)))
# print("token_mixer", x.shape)
x = x + self.drop_path(
self.layer_scale_2.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
* self.mlp(self.norm2(x)))
# print("mlp", x.shape)
else:
# print("==============not use_layer_scale======================")
x = x + self.drop_path(self.token_mixer(self.norm1(x)))
# print("token_mixer", x.shape)
x = x + self.drop_path(self.mlp(self.norm2(x)))
# print("mlp", x.shape)
return x
def basic_blocks_for_poolformer(dim, index, layers,
pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, norm_layer=GroupNorm,
drop_rate=.0, drop_path_rate=0.,
use_layer_scale=True, layer_scale_init_value=1e-5):
"""
generate PoolFormer blocks for a stage
return: PoolFormer blocks
"""
blocks = []
for block_idx in range(layers[index]):
block_dpr = drop_path_rate * (
block_idx + sum(layers[:index])) / (sum(layers) - 1)
blocks.append(PoolFormerBlock(
dim, pool_size=pool_size, mlp_ratio=mlp_ratio,
act_layer=act_layer, norm_layer=norm_layer,
drop=drop_rate, drop_path=block_dpr,
use_layer_scale=use_layer_scale,
layer_scale_init_value=layer_scale_init_value,
))
blocks = nn.Sequential(*blocks)
return blocks
class PoolFormer(nn.Module):
"""
PoolFormer, the main class of our model
--layers: [x,x,x,x], number of blocks for the 4 stages
--embed_dims, --mlp_ratios, --pool_size: the embedding dims, mlp ratios and
pooling size for the 4 stages
--downsamples: flags to apply downsampling or not
--norm_layer, --act_layer: define the types of normalization and activation
--num_classes: number of classes for the image classification
--in_patch_size, --in_stride, --in_pad: specify the patch embedding
for the input image
--down_patch_size --down_stride --down_pad:
specify the downsample (patch embed.)
--fork_feat: whether output features of the 4 stages, for dense prediction
--init_cfg, --pretrained:
for mmdetection and mmsegmentation to load pretrained weights
"""
def __init__(self, layers, embed_dims=None,
mlp_ratios=None, downsamples=None,
pool_size=3,
norm_layer=GroupNorm, act_layer=nn.GELU,
num_classes=1000,
in_patch_size=7, in_stride=4, in_pad=2,
down_patch_size=3, down_stride=2, down_pad=1,
drop_rate=0., drop_path_rate=0.,
use_layer_scale=True, layer_scale_init_value=1e-5,
fork_feat=False,
init_cfg=None,
pretrained=None,
**kwargs):
super().__init__()
if not fork_feat:
self.num_classes = num_classes
self.fork_feat = fork_feat
self.patch_embed = PatchEmbed(
patch_size=in_patch_size, stride=in_stride, padding=in_pad,
in_chans=1, embed_dim=embed_dims[0])
# set the main block in network
network = []
for i in range(len(layers)):
stage = basic_blocks_for_poolformer(embed_dims[i], i, layers,
pool_size=pool_size, mlp_ratio=mlp_ratios[i],
act_layer=act_layer, norm_layer=norm_layer,
drop_rate=drop_rate,
drop_path_rate=drop_path_rate,
use_layer_scale=use_layer_scale,
layer_scale_init_value=layer_scale_init_value)
network.append(stage)
if i >= len(layers) - 1:
break
if downsamples[i] or embed_dims[i] != embed_dims[i+1]:
# downsampling between two stages
network.append(
PatchEmbed(
patch_size=down_patch_size, stride=down_stride,
padding=down_pad,
in_chans=embed_dims[i], embed_dim=embed_dims[i+1]
)
)
self.network = nn.ModuleList(network)
if self.fork_feat:
# add a norm layer for each output
self.out_indices = [0, 2, 4, 6]
for i_emb, i_layer in enumerate(self.out_indices):
if i_emb == 0 and os.environ.get('FORK_LAST3', None):
# TODO: more elegant way
"""For RetinaNet, `start_level=1`. The first norm layer will not used.
cmd: `FORK_LAST3=1 python -m torch.distributed.launch ...`
"""
layer = nn.Identity()
else:
layer = norm_layer(embed_dims[i_emb])
layer_name = f'norm{i_layer}'
self.add_module(layer_name, layer)
else:
# Classifier head
self.norm = norm_layer(embed_dims[-1])
self.head = nn.Linear(
embed_dims[-1], num_classes) if num_classes > 0 \
else nn.Identity()
self.apply(self.cls_init_weights)
self.init_cfg = copy.deepcopy(init_cfg)
# load pre-trained model
if self.fork_feat and (
self.init_cfg is not None or pretrained is not None):
self.init_weights()
# init for classification
def cls_init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
def get_classifier(self):
return self.head
def reset_classifier(self, num_classes):
self.num_classes = num_classes
self.head = nn.Linear(
self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
def forward_embeddings(self, x):
x = self.patch_embed(x)
return x
def forward_tokens(self, x):
outs = []
layer_output = []
# print(f"network:{self.network}")
# print(f"first of tokens:{x.shape}")
# print("==========================================")
for idx, block in enumerate(self.network):
# print(f"idx:{idx}, block:{block}")
x = block(x)
if self.fork_feat and idx in self.out_indices:
# print("use the norm")
norm_layer = getattr(self, f'norm{idx}')
x_out = norm_layer(x)
outs.append(x_out)
# print(f"idx : {idx}, x.shape : {x.shape}")
if idx % 2 == 0:
layer_output.append(x)
if self.fork_feat:
# output the features of four stages for dense prediction
return outs
# output only the features of last layer for image classification
return layer_output
def forward(self, x):
# print("input", x.shape)
# input embedding
x = self.forward_embeddings(x)
# print("forward_embeddings", x.shape)
# through backbone
layer_output_list = self.forward_tokens(x)
x = layer_output_list[-1]
# print("forward_tokens", ([tmp, tmp.shape] for tmp in x))
# 打印每个张量的形状
# for i, tensor in enumerate(layer_output_list):
# print(f"Tensor {i}: Shape {tensor.shape}")
if self.fork_feat:
# otuput features of four stages for dense prediction
return x
x = self.norm(x)
# print("norm", x.shape)
# print("mean[-3, -2, -1]", x.mean([-3, -2, -1]).shape)
cls_out = self.head(x.mean([-3, -2, -1]))
layer_output_list.append(cls_out)
# print("cls_out", cls_out.shape)
# for image classification
return layer_output_list
def poolformer_s12(pretrained=False, **kwargs):
"""
PoolFormer-S12 model, Params: 12M
--layers: [x,x,x,x], numbers of layers for the four stages
--embed_dims, --mlp_ratios:
embedding dims and mlp ratios for the four stages
--downsamples: flags to apply downsampling or not in four blocks
"""
layers = [2, 2, 6, 2]
embed_dims = [64, 128, 320, 512]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
**kwargs)
return model
def poolformer_s24(pretrained=False, **kwargs):
"""
PoolFormer-S24 model, Params: 21M
"""
layers = [4, 4, 12, 4]
embed_dims = [64, 128, 320, 512]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
**kwargs)
return model
def poolformer_s36(pretrained=False, **kwargs):
"""
PoolFormer-S36 model, Params: 31M
"""
layers = [6, 6, 18, 6]
embed_dims = [64, 128, 320, 512]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
layer_scale_init_value=1e-6,
**kwargs)
return model
def poolformer_m36(pretrained=False, **kwargs):
"""
PoolFormer-M36 model, Params: 56M
"""
layers = [6, 6, 18, 6]
embed_dims = [96, 192, 384, 768]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
layer_scale_init_value=1e-6,
**kwargs)
return model
def poolformer_m48(pretrained=False, **kwargs):
"""
PoolFormer-M48 model, Params: 73M
"""
layers = [8, 8, 24, 8]
embed_dims = [96, 192, 384, 768]
mlp_ratios = [4, 4, 4, 4]
downsamples = [True, True, True, True]
model = PoolFormer(
layers, embed_dims=embed_dims,
mlp_ratios=mlp_ratios, downsamples=downsamples,
layer_scale_init_value=1e-6,
**kwargs)
return model
if __name__ == "__main__":
a = torch.randn(4, 1, 96, 128, 96)
model = poolformer_s12(num_classes=400)
print("参数量", sum(p.numel() for p in model.parameters()))
layer_output1, layer_output2, layer_output3, layer_output4, layer_output5 = model(a)
print(layer_output1.shape)
print(layer_output2.shape)
print(layer_output3.shape)
print(layer_output4.shape)
print(layer_output5.shape)
"""
参数量 15453776
torch.Size([4, 64, 24, 32, 24])
torch.Size([4, 128, 12, 16, 12])
torch.Size([4, 320, 6, 8, 6])
torch.Size([4, 512, 3, 4, 3])
torch.Size([4, 400])
"""2 Resnet
2.1 代码实现
2.1.1 导包
python
import torch
from torch import nn
import torch.nn.functional as F2.1.2 网络模块代码
python
# 因为resnet的18、34层，和50层以上有不同的基础结构，所以要设置两种基本残差卷积
# 对应18层和34层的基础残差结构
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_channels, out_channels, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels
, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.downsample = downsample
def forward(self, x):
identity = x
if self.downsample is not None: # 如果下采样函数存在，则对残差分支进行下采样
identity = self.downsample(identity)
conv1_out = self.relu(self.bn1(self.conv1(x))) # 卷积、bn、激活
print("conv1_out", conv1_out.shape)
conv2_out = self.bn2(self.conv2(conv1_out)) # 卷积、bn，激活函数要加上残差边后再使用
print("conv2_out", conv2_out.shape)
out_add = conv2_out + identity # 卷积输出加上残差边
out = self.relu(out_add) # 激活
return out
class BasicBlock2(nn.Module):
expansion = 4 # 在50层以上的结构中，需要利用1*1卷积进行升维，升的倍数是4倍
def __init__(self, in_channels, out_channels, stride=1, downsample=None):
super(BasicBlock2, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1,
stride=1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
stride=stride, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1,
stride=1, bias=False)
self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
self.downsample = downsample # 残差下采样
def forward(self, x):
identity = x # 保留残差分支
if self.downsample is not None:
identity = self.downsample(identity)
# 连续的卷积操作
conv1_out = self.relu(self.bn1(self.conv1(x)))
# print("conv1_out", conv1_out.shape)
conv2_out = self.relu(self.bn2(self.conv2(conv1_out)))
# print("conv2_out", conv2_out.shape)
conv3_out = self.bn3(self.conv3(conv2_out))
out_add = conv3_out + identity # 残差相加
out = self.relu(out_add)
return out
class ResNet(nn.Module):
def __init__(self, block, block_num, num_classes=1000, include_top=True):
# include_top是为了搭建其他网络时使用，include_top=False时不会采用全连接层，只要卷积主干特征提取网络
super(ResNet, self).__init__()
self.include_top = include_top
self.in_channels = 64 # resnet会先经过一个初始卷积，卷积后特征层通道数都是64
# 初始一个下采样卷积
self.conv1 = nn.Conv2d(3, out_channels=self.in_channels, kernel_size=7, stride=2,
padding=3, bias=False)
# (x-k+2p+1)/s+1，padding=3，使得图片的输出尺寸刚好为原来一般
self.bn1 = nn.BatchNorm2d(self.in_channels)
self.relu = nn.ReLU()
# padding=1,使得输出尺寸为原来一半，在默认dilation=1时，maxpool计算公式(h+2*p-k)/s+1
self.max_pool = nn.MaxPool2d(3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, block_num[0])
# stride=2是因为，从第二次开始，第一次卷积会将图片进行下采样
self.layer2 = self._make_layer(block, 128, block_num[1], stride=2)
self.layer3 = self._make_layer(block, 256, block_num[2], stride=2)
self.layer4 = self._make_layer(block, 512, block_num[3], stride=2)
if self.include_top:
# 通过平均池化下采样，无论特征层的宽高如何，输出都是1*1宽高
self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # output size = (1, 1)
# 线性网络，得到分类
self.fc = nn.Linear(512 * block.expansion, num_classes)
# 初始化操作
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
def forward(self, x):
# print("input", x.shape)
x = self.relu(self.bn1(self.conv1(x))) # 经过一个初始卷积
# print("x1", x.shape)
x = self.max_pool(x)
# print("x2", x.shape)
x = self.layer1(x)
# print("x3", x.shape)
x = self.layer2(x)
# print("x4", x.shape)
x = self.layer3(x)
# print("x5", x.shape)
x = self.layer4(x)
# print("x6", x.shape)
if self.include_top:
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
# 生成一个层，一层有多个基本残差block，数量由block_num控制
def _make_layer(self, block, channels, block_num, stride=1):
# channels:残差结构中第一层的卷积核的通道数
downsample = None
if stride != 1 or self.in_channels != channels * block.expansion:
# 下采样，通道数变换并且通过stride调整尺寸
downsample = nn.Sequential(
nn.Conv2d(self.in_channels, channels * block.expansion, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(channels * block.expansion)
)
layers = []
# 第一层可能会下采样
layers.append(block(self.in_channels, channels, downsample=downsample, stride=stride))
# 若是对于52层以上的，经过当前layers第一个卷积层后，通道数会扩张为当前channels的4倍
self.in_channels = channels * block.expansion
for _ in range(1, block_num):
layers.append(block(self.in_channels, channels))
return nn.Sequential(*layers)
# 构建18层的resnet
def resnet18(num_classes=1000, include_top=True):
# 预训练权重链接 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, include_top=include_top)
# 构建34层的resnet
def resnet34(num_classes=1000, include_top=True):
# 预训练权重链接 https://download.pytorch.org/models/resnet34-333f7ec4.pth
return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
# 可用于构建50层、101层和152层的resnet
# 对于resnet50 输入列表为:[3,4,6,3]
# 对于resnet101 输入列表为:[3,4,23,3]
# 对于resnet152 输入列表为:[3,8,36,3]
def resnet50(num_classes=1000, include_top=True):
# 预训练权重链接 https://download.pytorch.org/models/resnet50-19c8e357.pth
return ResNet(BasicBlock2, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
def resnet101(num_classes=1000, include_top=True):
# 预训练权重链接 https://download.pytorch.org/models/resnet101-5d3b4d8f.pth
return ResNet(BasicBlock2, [3, 4, 23, 3], num_classes=num_classes, include_top=include_top)2.1.3 测试网络输入输出
1、不用预训练的模型
python
if __name__ == "__main__":
model = resnet50()
x = torch.randn(1, 3, 224, 224)
out = model(x)
print("x", x.shape)
print("out", out.shape)输出：
text
x torch.Size([1, 3, 224, 224])
out torch.Size([1, 1000])2、用预训练的模型
Resnet50模型下载链接：https://download.pytorch.org/models/resnet50-19c8e357.pth
python
if __name__ == "__main__":
model = resnet50()
model.load_state_dict(torch.load(r"/path/of/your/model"))
x = torch.randn(1, 3, 224, 224)
out = model(x)
print("x", x.shape)
print("out", out.shape)输出：
text
x torch.Size([1, 3, 224, 224])
out torch.Size([1, 1000])