File size: 3,011 Bytes
102ac67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import logging
import torch
import torch.nn as nn
from mmcv.runner.checkpoint import load_state_dict
from mmdet.models.builder import BACKBONES
from .vit import ViT, SimpleFeaturePyramid, partial
from .fpn import LastLevelMaxPool
@BACKBONES.register_module()
class EVA02(nn.Module):
def __init__(
self,
# args for ViT
img_size=1024,
real_img_size=(256, 704),
patch_size=16,
in_chans=3,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4*2/3,
qkv_bias=True,
drop_path_rate=0.0,
norm_layer=partial(nn.LayerNorm, eps=1e-6),
use_abs_pos=True,
pt_hw_seq_len=16,
intp_freq=True,
window_size=0,
window_block_indexes=(),
residual_block_indexes=(),
use_act_checkpoint=False,
pretrain_img_size=224,
pretrain_use_cls_token=True,
out_feature="last_feat",
xattn=False,
frozen_blocks=-1,
# args for simple FPN
fpn_in_feature="last_feat",
fpn_out_channels=256,
fpn_scale_factors=(4.0, 2.0, 1.0, 0.5),
fpn_top_block=False,
fpn_norm="LN",
fpn_square_pad=0,
pretrained=None
):
super().__init__()
self.backbone = SimpleFeaturePyramid(
ViT(
img_size=img_size,
real_img_size=real_img_size,
patch_size=patch_size,
in_chans=in_chans,
embed_dim=embed_dim,
depth=depth,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
drop_path_rate=drop_path_rate,
norm_layer=norm_layer,
use_abs_pos=use_abs_pos,
pt_hw_seq_len=pt_hw_seq_len,
intp_freq=intp_freq,
window_size=window_size,
window_block_indexes=window_block_indexes,
residual_block_indexes=residual_block_indexes,
use_act_checkpoint=use_act_checkpoint,
pretrain_img_size=pretrain_img_size,
pretrain_use_cls_token=pretrain_use_cls_token,
out_feature=out_feature,
xattn=xattn,
frozen_blocks=frozen_blocks,
),
in_feature=fpn_in_feature,
out_channels=fpn_out_channels,
scale_factors=fpn_scale_factors,
top_block=LastLevelMaxPool() if fpn_top_block else None,
norm=fpn_norm,
square_pad=fpn_square_pad,
)
self.init_weights(pretrained)
def init_weights(self, pretrained=None):
if pretrained is None:
return
logging.info('Loading pretrained weights from %s' % pretrained)
state_dict = torch.load(pretrained)['model']
load_state_dict(self, state_dict, strict=False)
def forward(self, x):
outs = self.backbone(x)
return list(outs.values())
|