File size: 3,283 Bytes
da2e2ac |
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 |
"""
Implements the TransFuser vision backbone.
"""
import timm
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.checkpoint import checkpoint
from navsim.agents.backbones.internimage import InternImage
from navsim.agents.backbones.swin import SwinTransformerBEVFT
from navsim.agents.backbones.vov import VoVNet
from navsim.agents.dreamer.hydra_dreamer_config import HydraDreamerConfig
from navsim.agents.hydra.hydra_config import HydraConfig
from navsim.agents.transfuser.transfuser_backbone import GPT
from navsim.agents.utils.vit import DAViT
class Backbone(nn.Module):
"""
Multi-scale Fusion Transformer for image + LiDAR feature fusion
"""
def __init__(self, config: HydraDreamerConfig):
super().__init__()
self.config = config
self.backbone_type = config.backbone_type
if config.backbone_type == 'intern':
self.image_encoder = InternImage(init_cfg=dict(type='Pretrained',
checkpoint=config.intern_ckpt
),
frozen_stages=2)
# scale_4_c = 2560
vit_channels = 2560
self.image_encoder.init_weights()
elif config.backbone_type == 'vov':
self.image_encoder = VoVNet(
spec_name='V-99-eSE',
out_features=['stage4', 'stage5'],
norm_eval=True,
with_cp=True,
init_cfg=dict(
type='Pretrained',
checkpoint=config.vov_ckpt,
prefix='img_backbone.'
)
)
# scale_4_c = 1024
vit_channels = 1024
self.image_encoder.init_weights()
elif config.backbone_type == 'swin':
self.image_encoder = SwinTransformerBEVFT(
with_cp=True,
convert_weights=False,
depths=[2,2,18,2],
drop_path_rate=0.35,
embed_dims=192,
init_cfg=dict(
checkpoint=config.swin_ckpt,
type='Pretrained'
),
num_heads=[6,12,24,48],
out_indices=[3],
patch_norm=True,
window_size=[16,16,16,16],
use_abs_pos_embed=True,
return_stereo_feat=False,
output_missing_index_as_none=False
)
vit_channels = 1536
elif config.backbone_type == 'vit':
self.image_encoder = DAViT(ckpt=config.vit_ckpt)
vit_channels = 1024
elif config.backbone_type == 'resnet':
self.image_encoder = timm.create_model(
'resnet34', pretrained=False, features_only=True
)
vit_channels = 512
else:
raise ValueError
self.avgpool_img = nn.AdaptiveAvgPool2d(
(self.config.img_vert_anchors, self.config.img_horz_anchors)
)
self.img_feat_c = vit_channels
def forward(self, image):
image_features = self.image_encoder(image)[-1]
return self.avgpool_img(image_features)
|