| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589 |
- # copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- # Code was based on https://github.com/facebookresearch/LeViT
- import itertools
- import math
- import warnings
- import paddle
- import paddle.nn as nn
- import paddle.nn.functional as F
- from paddle.nn.initializer import TruncatedNormal, Constant
- from paddle.regularizer import L2Decay
- from .vision_transformer import trunc_normal_, zeros_, ones_, Identity
- from paddlex.ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
- MODEL_URLS = {
- "LeViT_128S":
- "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/LeViT_128S_pretrained.pdparams",
- "LeViT_128":
- "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/LeViT_128_pretrained.pdparams",
- "LeViT_192":
- "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/LeViT_192_pretrained.pdparams",
- "LeViT_256":
- "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/LeViT_256_pretrained.pdparams",
- "LeViT_384":
- "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/LeViT_384_pretrained.pdparams",
- }
- __all__ = list(MODEL_URLS.keys())
- def cal_attention_biases(attention_biases, attention_bias_idxs):
- gather_list = []
- attention_bias_t = paddle.transpose(attention_biases, (1, 0))
- nums = attention_bias_idxs.shape[0]
- for idx in range(nums):
- gather = paddle.gather(attention_bias_t, attention_bias_idxs[idx])
- gather_list.append(gather)
- shape0, shape1 = attention_bias_idxs.shape
- gather = paddle.concat(gather_list)
- return paddle.transpose(gather, (1, 0)).reshape((0, shape0, shape1))
- class Conv2d_BN(nn.Sequential):
- def __init__(self,
- a,
- b,
- ks=1,
- stride=1,
- pad=0,
- dilation=1,
- groups=1,
- bn_weight_init=1,
- resolution=-10000):
- super().__init__()
- self.add_sublayer(
- 'c',
- nn.Conv2D(
- a, b, ks, stride, pad, dilation, groups, bias_attr=False))
- bn = nn.BatchNorm2D(b)
- ones_(bn.weight)
- zeros_(bn.bias)
- self.add_sublayer('bn', bn)
- class Linear_BN(nn.Sequential):
- def __init__(self, a, b, bn_weight_init=1):
- super().__init__()
- self.add_sublayer('c', nn.Linear(a, b, bias_attr=False))
- bn = nn.BatchNorm1D(b)
- if bn_weight_init == 0:
- zeros_(bn.weight)
- else:
- ones_(bn.weight)
- zeros_(bn.bias)
- self.add_sublayer('bn', bn)
- def forward(self, x):
- l, bn = self._sub_layers.values()
- x = l(x)
- return paddle.reshape(bn(x.flatten(0, 1)), x.shape)
- class BN_Linear(nn.Sequential):
- def __init__(self, a, b, bias=True, std=0.02):
- super().__init__()
- self.add_sublayer('bn', nn.BatchNorm1D(a))
- l = nn.Linear(a, b, bias_attr=bias)
- trunc_normal_(l.weight)
- if bias:
- zeros_(l.bias)
- self.add_sublayer('l', l)
- def b16(n, activation, resolution=224):
- return nn.Sequential(
- Conv2d_BN(
- 3, n // 8, 3, 2, 1, resolution=resolution),
- activation(),
- Conv2d_BN(
- n // 8, n // 4, 3, 2, 1, resolution=resolution // 2),
- activation(),
- Conv2d_BN(
- n // 4, n // 2, 3, 2, 1, resolution=resolution // 4),
- activation(),
- Conv2d_BN(
- n // 2, n, 3, 2, 1, resolution=resolution // 8))
- class Residual(nn.Layer):
- def __init__(self, m, drop):
- super().__init__()
- self.m = m
- self.drop = drop
- def forward(self, x):
- if self.training and self.drop > 0:
- y = paddle.rand(
- shape=[x.shape[0], 1, 1]).__ge__(self.drop).astype("float32")
- y = y.divide(paddle.full_like(y, 1 - self.drop))
- return paddle.add(x, y)
- else:
- return paddle.add(x, self.m(x))
- class Attention(nn.Layer):
- def __init__(self,
- dim,
- key_dim,
- num_heads=8,
- attn_ratio=4,
- activation=None,
- resolution=14):
- super().__init__()
- self.num_heads = num_heads
- self.scale = key_dim**-0.5
- self.key_dim = key_dim
- self.nh_kd = nh_kd = key_dim * num_heads
- self.d = int(attn_ratio * key_dim)
- self.dh = int(attn_ratio * key_dim) * num_heads
- self.attn_ratio = attn_ratio
- self.h = self.dh + nh_kd * 2
- self.qkv = Linear_BN(dim, self.h)
- self.proj = nn.Sequential(
- activation(), Linear_BN(
- self.dh, dim, bn_weight_init=0))
- points = list(itertools.product(range(resolution), range(resolution)))
- N = len(points)
- attention_offsets = {}
- idxs = []
- for p1 in points:
- for p2 in points:
- offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
- if offset not in attention_offsets:
- attention_offsets[offset] = len(attention_offsets)
- idxs.append(attention_offsets[offset])
- self.attention_biases = self.create_parameter(
- shape=(num_heads, len(attention_offsets)),
- default_initializer=zeros_,
- attr=paddle.ParamAttr(regularizer=L2Decay(0.0)))
- tensor_idxs = paddle.to_tensor(idxs, dtype='int64')
- self.register_buffer('attention_bias_idxs',
- paddle.reshape(tensor_idxs, [N, N]))
- @paddle.no_grad()
- def train(self, mode=True):
- if mode:
- super().train()
- else:
- super().eval()
- if mode and hasattr(self, 'ab'):
- del self.ab
- else:
- self.ab = cal_attention_biases(self.attention_biases,
- self.attention_bias_idxs)
- def forward(self, x):
- self.training = True
- B, N, C = x.shape
- qkv = self.qkv(x)
- qkv = paddle.reshape(qkv,
- [B, N, self.num_heads, self.h // self.num_heads])
- q, k, v = paddle.split(
- qkv, [self.key_dim, self.key_dim, self.d], axis=3)
- q = paddle.transpose(q, perm=[0, 2, 1, 3])
- k = paddle.transpose(k, perm=[0, 2, 1, 3])
- v = paddle.transpose(v, perm=[0, 2, 1, 3])
- k_transpose = paddle.transpose(k, perm=[0, 1, 3, 2])
- if self.training:
- attention_biases = cal_attention_biases(self.attention_biases,
- self.attention_bias_idxs)
- else:
- attention_biases = self.ab
- attn = (paddle.matmul(q, k_transpose) * self.scale + attention_biases)
- attn = F.softmax(attn)
- x = paddle.transpose(paddle.matmul(attn, v), perm=[0, 2, 1, 3])
- x = paddle.reshape(x, [B, N, self.dh])
- x = self.proj(x)
- return x
- class Subsample(nn.Layer):
- def __init__(self, stride, resolution):
- super().__init__()
- self.stride = stride
- self.resolution = resolution
- def forward(self, x):
- B, N, C = x.shape
- x = paddle.reshape(x, [B, self.resolution, self.resolution, C])
- end1, end2 = x.shape[1], x.shape[2]
- x = x[:, 0:end1:self.stride, 0:end2:self.stride]
- x = paddle.reshape(x, [B, -1, C])
- return x
- class AttentionSubsample(nn.Layer):
- def __init__(self,
- in_dim,
- out_dim,
- key_dim,
- num_heads=8,
- attn_ratio=2,
- activation=None,
- stride=2,
- resolution=14,
- resolution_=7):
- super().__init__()
- self.num_heads = num_heads
- self.scale = key_dim**-0.5
- self.key_dim = key_dim
- self.nh_kd = nh_kd = key_dim * num_heads
- self.d = int(attn_ratio * key_dim)
- self.dh = int(attn_ratio * key_dim) * self.num_heads
- self.attn_ratio = attn_ratio
- self.resolution_ = resolution_
- self.resolution_2 = resolution_**2
- self.training = True
- h = self.dh + nh_kd
- self.kv = Linear_BN(in_dim, h)
- self.q = nn.Sequential(
- Subsample(stride, resolution), Linear_BN(in_dim, nh_kd))
- self.proj = nn.Sequential(activation(), Linear_BN(self.dh, out_dim))
- self.stride = stride
- self.resolution = resolution
- points = list(itertools.product(range(resolution), range(resolution)))
- points_ = list(
- itertools.product(range(resolution_), range(resolution_)))
- N = len(points)
- N_ = len(points_)
- attention_offsets = {}
- idxs = []
- i = 0
- j = 0
- for p1 in points_:
- i += 1
- for p2 in points:
- j += 1
- size = 1
- offset = (abs(p1[0] * stride - p2[0] + (size - 1) / 2),
- abs(p1[1] * stride - p2[1] + (size - 1) / 2))
- if offset not in attention_offsets:
- attention_offsets[offset] = len(attention_offsets)
- idxs.append(attention_offsets[offset])
- self.attention_biases = self.create_parameter(
- shape=(num_heads, len(attention_offsets)),
- default_initializer=zeros_,
- attr=paddle.ParamAttr(regularizer=L2Decay(0.0)))
- tensor_idxs_ = paddle.to_tensor(idxs, dtype='int64')
- self.register_buffer('attention_bias_idxs',
- paddle.reshape(tensor_idxs_, [N_, N]))
- @paddle.no_grad()
- def train(self, mode=True):
- if mode:
- super().train()
- else:
- super().eval()
- if mode and hasattr(self, 'ab'):
- del self.ab
- else:
- self.ab = cal_attention_biases(self.attention_biases,
- self.attention_bias_idxs)
- def forward(self, x):
- self.training = True
- B, N, C = x.shape
- kv = self.kv(x)
- kv = paddle.reshape(kv, [B, N, self.num_heads, -1])
- k, v = paddle.split(kv, [self.key_dim, self.d], axis=3)
- k = paddle.transpose(k, perm=[0, 2, 1, 3]) # BHNC
- v = paddle.transpose(v, perm=[0, 2, 1, 3])
- q = paddle.reshape(
- self.q(x), [B, self.resolution_2, self.num_heads, self.key_dim])
- q = paddle.transpose(q, perm=[0, 2, 1, 3])
- if self.training:
- attention_biases = cal_attention_biases(self.attention_biases,
- self.attention_bias_idxs)
- else:
- attention_biases = self.ab
- attn = (paddle.matmul(
- q, paddle.transpose(
- k, perm=[0, 1, 3, 2]))) * self.scale + attention_biases
- attn = F.softmax(attn)
- x = paddle.reshape(
- paddle.transpose(
- paddle.matmul(attn, v), perm=[0, 2, 1, 3]), [B, -1, self.dh])
- x = self.proj(x)
- return x
- class LeViT(nn.Layer):
- """ Vision Transformer with support for patch or hybrid CNN input stage
- """
- def __init__(self,
- img_size=224,
- patch_size=16,
- in_chans=3,
- class_num=1000,
- embed_dim=[192],
- key_dim=[64],
- depth=[12],
- num_heads=[3],
- attn_ratio=[2],
- mlp_ratio=[2],
- hybrid_backbone=None,
- down_ops=[],
- attention_activation=nn.Hardswish,
- mlp_activation=nn.Hardswish,
- distillation=True,
- drop_path=0):
- super().__init__()
- self.class_num = class_num
- self.num_features = embed_dim[-1]
- self.embed_dim = embed_dim
- self.distillation = distillation
- self.patch_embed = hybrid_backbone
- self.blocks = []
- down_ops.append([''])
- resolution = img_size // patch_size
- for i, (ed, kd, dpth, nh, ar, mr, do) in enumerate(
- zip(embed_dim, key_dim, depth, num_heads, attn_ratio,
- mlp_ratio, down_ops)):
- for _ in range(dpth):
- self.blocks.append(
- Residual(
- Attention(
- ed,
- kd,
- nh,
- attn_ratio=ar,
- activation=attention_activation,
- resolution=resolution, ),
- drop_path))
- if mr > 0:
- h = int(ed * mr)
- self.blocks.append(
- Residual(
- nn.Sequential(
- Linear_BN(ed, h),
- mlp_activation(),
- Linear_BN(
- h, ed, bn_weight_init=0), ),
- drop_path))
- if do[0] == 'Subsample':
- #('Subsample',key_dim, num_heads, attn_ratio, mlp_ratio, stride)
- resolution_ = (resolution - 1) // do[5] + 1
- self.blocks.append(
- AttentionSubsample(
- *embed_dim[i:i + 2],
- key_dim=do[1],
- num_heads=do[2],
- attn_ratio=do[3],
- activation=attention_activation,
- stride=do[5],
- resolution=resolution,
- resolution_=resolution_))
- resolution = resolution_
- if do[4] > 0: # mlp_ratio
- h = int(embed_dim[i + 1] * do[4])
- self.blocks.append(
- Residual(
- nn.Sequential(
- Linear_BN(embed_dim[i + 1], h),
- mlp_activation(),
- Linear_BN(
- h, embed_dim[i + 1], bn_weight_init=0), ),
- drop_path))
- self.blocks = nn.Sequential(*self.blocks)
- # Classifier head
- self.head = BN_Linear(embed_dim[-1],
- class_num) if class_num > 0 else Identity()
- if distillation:
- self.head_dist = BN_Linear(
- embed_dim[-1], class_num) if class_num > 0 else Identity()
- def forward(self, x):
- x = self.patch_embed(x)
- x = x.flatten(2)
- x = paddle.transpose(x, perm=[0, 2, 1])
- x = self.blocks(x)
- x = x.mean(1)
- x = paddle.reshape(x, [-1, self.embed_dim[-1]])
- if self.distillation:
- x = self.head(x), self.head_dist(x)
- if not self.training:
- x = (x[0] + x[1]) / 2
- else:
- x = self.head(x)
- return x
- def model_factory(C, D, X, N, drop_path, class_num, distillation):
- embed_dim = [int(x) for x in C.split('_')]
- num_heads = [int(x) for x in N.split('_')]
- depth = [int(x) for x in X.split('_')]
- act = nn.Hardswish
- model = LeViT(
- patch_size=16,
- embed_dim=embed_dim,
- num_heads=num_heads,
- key_dim=[D] * 3,
- depth=depth,
- attn_ratio=[2, 2, 2],
- mlp_ratio=[2, 2, 2],
- down_ops=[
- #('Subsample',key_dim, num_heads, attn_ratio, mlp_ratio, stride)
- ['Subsample', D, embed_dim[0] // D, 4, 2, 2],
- ['Subsample', D, embed_dim[1] // D, 4, 2, 2],
- ],
- attention_activation=act,
- mlp_activation=act,
- hybrid_backbone=b16(embed_dim[0], activation=act),
- class_num=class_num,
- drop_path=drop_path,
- distillation=distillation)
- return model
- specification = {
- 'LeViT_128S': {
- 'C': '128_256_384',
- 'D': 16,
- 'N': '4_6_8',
- 'X': '2_3_4',
- 'drop_path': 0
- },
- 'LeViT_128': {
- 'C': '128_256_384',
- 'D': 16,
- 'N': '4_8_12',
- 'X': '4_4_4',
- 'drop_path': 0
- },
- 'LeViT_192': {
- 'C': '192_288_384',
- 'D': 32,
- 'N': '3_5_6',
- 'X': '4_4_4',
- 'drop_path': 0
- },
- 'LeViT_256': {
- 'C': '256_384_512',
- 'D': 32,
- 'N': '4_6_8',
- 'X': '4_4_4',
- 'drop_path': 0
- },
- 'LeViT_384': {
- 'C': '384_512_768',
- 'D': 32,
- 'N': '6_9_12',
- 'X': '4_4_4',
- 'drop_path': 0.1
- },
- }
- def _load_pretrained(pretrained, model, model_url, use_ssld=False):
- if pretrained is False:
- pass
- elif pretrained is True:
- load_dygraph_pretrain_from_url(model, model_url, use_ssld=use_ssld)
- elif isinstance(pretrained, str):
- load_dygraph_pretrain(model, pretrained)
- else:
- raise RuntimeError(
- "pretrained type is not available. Please use `string` or `boolean` type."
- )
- def LeViT_128S(pretrained=False,
- use_ssld=False,
- class_num=1000,
- distillation=False,
- **kwargs):
- model = model_factory(
- **specification['LeViT_128S'],
- class_num=class_num,
- distillation=distillation)
- _load_pretrained(
- pretrained, model, MODEL_URLS["LeViT_128S"], use_ssld=use_ssld)
- return model
- def LeViT_128(pretrained=False,
- use_ssld=False,
- class_num=1000,
- distillation=False,
- **kwargs):
- model = model_factory(
- **specification['LeViT_128'],
- class_num=class_num,
- distillation=distillation)
- _load_pretrained(
- pretrained, model, MODEL_URLS["LeViT_128"], use_ssld=use_ssld)
- return model
- def LeViT_192(pretrained=False,
- use_ssld=False,
- class_num=1000,
- distillation=False,
- **kwargs):
- model = model_factory(
- **specification['LeViT_192'],
- class_num=class_num,
- distillation=distillation)
- _load_pretrained(
- pretrained, model, MODEL_URLS["LeViT_192"], use_ssld=use_ssld)
- return model
- def LeViT_256(pretrained=False,
- use_ssld=False,
- class_num=1000,
- distillation=False,
- **kwargs):
- model = model_factory(
- **specification['LeViT_256'],
- class_num=class_num,
- distillation=distillation)
- _load_pretrained(
- pretrained, model, MODEL_URLS["LeViT_256"], use_ssld=use_ssld)
- return model
- def LeViT_384(pretrained=False,
- use_ssld=False,
- class_num=1000,
- distillation=False,
- **kwargs):
- model = model_factory(
- **specification['LeViT_384'],
- class_num=class_num,
- distillation=distillation)
- _load_pretrained(
- pretrained, model, MODEL_URLS["LeViT_384"], use_ssld=use_ssld)
- return model
|