Upload folder using huggingface_hub

model.py

@@ -619,456 +619,95 @@ def build_attn_mask(mask_type):
         mask[:, 101:151] = False
     return mask
 
-class Interfuser(nn.Module):
+class InterfuserConfig(PretrainedConfig):
+    model_type = "interfuser"
+
     def __init__(
         self,
-        img_size=224,
-        multi_view_img_size=112,
-        patch_size=8,
-        in_chans=3,
-        embed_dim=768,
+        embed_dim=256,
         enc_depth=6,
         dec_depth=6,
-        dim_feedforward=2048,
-        normalize_before=False,
-        rgb_backbone_name="r26",
-        lidar_backbone_name="r26",
         num_heads=8,
-        norm_layer=None,
+        dim_feedforward=2048,
         dropout=0.1,
-        end2end=False,
-        direct_concat=True,
-        separate_view_attention=False,
-        separate_all_attention=False,
-        act_layer=None,
-        weight_init="",
-        freeze_num=-1,
-        with_lidar=False,
-        with_right_left_sensors=True,
-        with_center_sensor=False,
-        traffic_pred_head_type="det",
-        waypoints_pred_head="heatmap",
-        reverse_pos=True,
-        use_different_backbone=False,
-        use_view_embed=True,
-        use_mmad_pretrain=None,
+        rgb_backbone_name="r50",
+        lidar_backbone_name="r18",
+        use_different_backbone=True,
+        waypoints_pred_head="gru",
+        **kwargs,
     ):
-        super().__init__()
-        self.traffic_pred_head_type = traffic_pred_head_type
-        self.num_features = (
-            self.embed_dim
-        ) = embed_dim  # num_features for consistency with other models
-        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
-        act_layer = act_layer or nn.GELU
-
-        self.reverse_pos = reverse_pos
+        self.embed_dim = embed_dim
+        self.enc_depth = enc_depth
+        self.dec_depth = dec_depth
+        self.num_heads = num_heads
+        self.dim_feedforward = dim_feedforward
+        self.dropout = dropout
+        self.rgb_backbone_name = rgb_backbone_name
+        self.lidar_backbone_name = lidar_backbone_name
+        self.use_different_backbone = use_different_backbone
         self.waypoints_pred_head = waypoints_pred_head
-        self.with_lidar = with_lidar
-        self.with_right_left_sensors = with_right_left_sensors
-        self.with_center_sensor = with_center_sensor
-
-        self.direct_concat = direct_concat
-        self.separate_view_attention = separate_view_attention
-        self.separate_all_attention = separate_all_attention
-        self.end2end = end2end
-        self.use_view_embed = use_view_embed
-
-        if self.direct_concat:
-            in_chans = in_chans * 4
-            self.with_center_sensor = False
-            self.with_right_left_sensors = False
-
-        if self.separate_view_attention:
-            self.attn_mask = build_attn_mask("seperate_view")
-        elif self.separate_all_attention:
-            self.attn_mask = build_attn_mask("seperate_all")
-        else:
-            self.attn_mask = None
-
+        super().__init__(**kwargs)
+
+class InterfuserModel(PreTrainedModel):
+    config_class = InterfuserConfig
+
+    def __init__(self, config: InterfuserConfig):
+        super().__init__(config)
+        self.config = config
+        
+        embed_dim=config.embed_dim; enc_depth=config.enc_depth; dec_depth=config.dec_depth; num_heads=config.num_heads; dim_feedforward=config.dim_feedforward; dropout=config.dropout
+        rgb_backbone_name=config.rgb_backbone_name; lidar_backbone_name=config.lidar_backbone_name; use_different_backbone=config.use_different_backbone
+        in_chans=3; img_size=224; direct_concat=True; with_lidar=True; with_right_left_sensors=True
+
+        self.embed_dim = embed_dim; self.direct_concat = direct_concat; self.with_lidar = with_lidar; self.with_right_left_sensors = with_right_left_sensors
+        in_chans_rgb = in_chans * 4 if self.direct_concat else in_chans
+        
         if use_different_backbone:
-            if rgb_backbone_name == "r50":
-                self.rgb_backbone = resnet50d(
-                    pretrained=True,
-                    in_chans=in_chans,
-                    features_only=True,
-                    out_indices=[4],
-                )
-            elif rgb_backbone_name == "r26":
-                self.rgb_backbone = resnet26d(
-                    pretrained=True,
-                    in_chans=in_chans,
-                    features_only=True,
-                    out_indices=[4],
-                )
-            elif rgb_backbone_name == "r18":
-                self.rgb_backbone = resnet18d(
-                    pretrained=True,
-                    in_chans=in_chans,
-                    features_only=True,
-                    out_indices=[4],
-                )
-            if lidar_backbone_name == "r50":
-                self.lidar_backbone = resnet50d(
-                    pretrained=False,
-                    in_chans=in_chans,
-                    features_only=True,
-                    out_indices=[4],
-                )
-            elif lidar_backbone_name == "r26":
-                self.lidar_backbone = resnet26d(
-                    pretrained=False,
-                    in_chans=in_chans,
-                    features_only=True,
-                    out_indices=[4],
-                )
-            elif lidar_backbone_name == "r18":
-                self.lidar_backbone = resnet18d(
-                    pretrained=False, in_chans=3, features_only=True, out_indices=[4]
-                )
-            rgb_embed_layer = partial(HybridEmbed, backbone=self.rgb_backbone)
-            lidar_embed_layer = partial(HybridEmbed, backbone=self.lidar_backbone)
-
-            if use_mmad_pretrain:
-                params = torch.load(use_mmad_pretrain)["state_dict"]
-                updated_params = OrderedDict()
-                for key in params:
-                    if "backbone" in key:
-                        updated_params[key.replace("backbone.", "")] = params[key]
-                self.rgb_backbone.load_state_dict(updated_params)
-
-            self.rgb_patch_embed = rgb_embed_layer(
-                img_size=img_size,
-                patch_size=patch_size,
-                in_chans=in_chans,
-                embed_dim=embed_dim,
-            )
-            self.lidar_patch_embed = lidar_embed_layer(
-                img_size=img_size,
-                patch_size=patch_size,
-                in_chans=3,
-                embed_dim=embed_dim,
-            )
-        else:
-            if rgb_backbone_name == "r50":
-                self.rgb_backbone = resnet50d(
-                    pretrained=True, in_chans=3, features_only=True, out_indices=[4]
-                )
-            elif rgb_backbone_name == "r101":
-                self.rgb_backbone = resnet101d(
-                    pretrained=True, in_chans=3, features_only=True, out_indices=[4]
-                )
-            elif rgb_backbone_name == "r26":
-                self.rgb_backbone = resnet26d(
-                    pretrained=True, in_chans=3, features_only=True, out_indices=[4]
-                )
-            elif rgb_backbone_name == "r18":
-                self.rgb_backbone = resnet18d(
-                    pretrained=True, in_chans=3, features_only=True, out_indices=[4]
-                )
-            embed_layer = partial(HybridEmbed, backbone=self.rgb_backbone)
-
-            self.rgb_patch_embed = embed_layer(
-                img_size=img_size,
-                patch_size=patch_size,
-                in_chans=in_chans,
-                embed_dim=embed_dim,
-            )
-            self.lidar_patch_embed = embed_layer(
-                img_size=img_size,
-                patch_size=patch_size,
-                in_chans=in_chans,
-                embed_dim=embed_dim,
-            )
-
-        self.global_embed = nn.Parameter(torch.zeros(1, embed_dim, 5))
-        self.view_embed = nn.Parameter(torch.zeros(1, embed_dim, 5, 1))
-
-        if self.end2end:
-            self.query_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, 4))
-            self.query_embed = nn.Parameter(torch.zeros(4, 1, embed_dim))
-        elif self.waypoints_pred_head == "heatmap":
-            self.query_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, 5))
-            self.query_embed = nn.Parameter(torch.zeros(400 + 5, 1, embed_dim))
-        else:
-            self.query_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, 11))
-            self.query_embed = nn.Parameter(torch.zeros(400 + 11, 1, embed_dim))
-
-        if self.end2end:
-            self.waypoints_generator = GRUWaypointsPredictor(embed_dim, 4)
-        elif self.waypoints_pred_head == "heatmap":
-            self.waypoints_generator = MultiPath_Generator(
-                embed_dim + 32, embed_dim, 10
-            )
-        elif self.waypoints_pred_head == "gru":
-            self.waypoints_generator = GRUWaypointsPredictor(embed_dim)
-        elif self.waypoints_pred_head == "gru-command":
-            self.waypoints_generator = GRUWaypointsPredictorWithCommand(embed_dim)
-        elif self.waypoints_pred_head == "linear":
-            self.waypoints_generator = LinearWaypointsPredictor(embed_dim)
-        elif self.waypoints_pred_head == "linear-sum":
-            self.waypoints_generator = LinearWaypointsPredictor(embed_dim, cumsum=True)
-
-        self.junction_pred_head = nn.Linear(embed_dim, 2)
-        self.traffic_light_pred_head = nn.Linear(embed_dim, 2)
-        self.stop_sign_head = nn.Linear(embed_dim, 2)
-
-        if self.traffic_pred_head_type == "det":
-            self.traffic_pred_head = nn.Sequential(
-                *[
-                    nn.Linear(embed_dim + 32, 64),
-                    nn.ReLU(),
-                    nn.Linear(64, 7),
-                    nn.Sigmoid(),
-                ]
-            )
-        elif self.traffic_pred_head_type == "seg":
-            self.traffic_pred_head = nn.Sequential(
-                *[nn.Linear(embed_dim, 64), nn.ReLU(), nn.Linear(64, 1), nn.Sigmoid()]
-            )
-
+            self.rgb_backbone = {'r50': resnet50d, 'r26': resnet26d, 'r18': resnet18d}[rgb_backbone_name](pretrained=False, in_chans=in_chans_rgb, features_only=True, out_indices=[4])
+            self.lidar_backbone = {'r50': resnet50d, 'r26': resnet26d, 'r18': resnet18d}[lidar_backbone_name](pretrained=False, in_chans=in_chans, features_only=True, out_indices=[4])
+            self.rgb_patch_embed = HybridEmbed(self.rgb_backbone, img_size=img_size, in_chans=in_chans_rgb, embed_dim=embed_dim)
+            self.lidar_patch_embed = HybridEmbed(self.lidar_backbone, img_size=112, in_chans=in_chans, embed_dim=embed_dim)
+        
+        self.global_embed = nn.Parameter(torch.zeros(1, embed_dim, 5)); self.view_embed = nn.Parameter(torch.zeros(1, embed_dim, 5, 1))
+        self.query_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, 11)); self.query_embed = nn.Parameter(torch.zeros(400 + 11, 1, embed_dim))
+        self.waypoints_generator = GRUWaypointsPredictor(embed_dim)
+        self.junction_pred_head = nn.Linear(embed_dim, 2); self.traffic_light_pred_head = nn.Linear(embed_dim, 2); self.stop_sign_head = nn.Linear(embed_dim, 2)
+        self.traffic_pred_head = nn.Sequential(*[nn.Linear(embed_dim, 64), nn.ReLU(), nn.Linear(64, 7), nn.Sigmoid()])
         self.position_encoding = PositionEmbeddingSine(embed_dim // 2, normalize=True)
-
-        encoder_layer = TransformerEncoderLayer(
-            embed_dim, num_heads, dim_feedforward, dropout, act_layer, normalize_before
-        )
+        act_layer = nn.GELU()
+        encoder_layer = TransformerEncoderLayer(embed_dim, num_heads, dim_feedforward, dropout, act_layer)
         self.encoder = TransformerEncoder(encoder_layer, enc_depth, None)
+        decoder_layer = TransformerDecoderLayer(embed_dim, num_heads, dim_feedforward, dropout, act_layer)
+        self.decoder = TransformerDecoder(decoder_layer, dec_depth, nn.LayerNorm(embed_dim), return_intermediate=False)
 
-        decoder_layer = TransformerDecoderLayer(
-            embed_dim, num_heads, dim_feedforward, dropout, act_layer, normalize_before
-        )
-        decoder_norm = nn.LayerNorm(embed_dim)
-        self.decoder = TransformerDecoder(
-            decoder_layer, dec_depth, decoder_norm, return_intermediate=False
-        )
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        nn.init.uniform_(self.global_embed)
-        nn.init.uniform_(self.view_embed)
-        nn.init.uniform_(self.query_embed)
-        nn.init.uniform_(self.query_pos_embed)
-
-    def forward_features(
-        self,
-        front_image,
-        left_image,
-        right_image,
-        front_center_image,
-        lidar,
-        measurements,
-    ):
-        features = []
-
-        # Front view processing
-        front_image_token, front_image_token_global = self.rgb_patch_embed(front_image)
-        if self.use_view_embed:
-            front_image_token = (
-                front_image_token
-                + self.view_embed[:, :, 0:1, :]
-                + self.position_encoding(front_image_token)
-            )
-        else:
-            front_image_token = front_image_token + self.position_encoding(
-                front_image_token
-            )
-        front_image_token = front_image_token.flatten(2).permute(2, 0, 1)
-        front_image_token_global = (
-            front_image_token_global
-            + self.view_embed[:, :, 0, :]
-            + self.global_embed[:, :, 0:1]
-        )
-        front_image_token_global = front_image_token_global.permute(2, 0, 1)
-        features.extend([front_image_token, front_image_token_global])
-
+    def forward_features(self, front_image, left_image, right_image, lidar):
+        features = [];
+        x, x_g = self.rgb_patch_embed(front_image); x = x + self.view_embed[:,:,0:1,:] + self.position_encoding(x); x=x.flatten(2).permute(2,0,1); x_g=x_g+self.view_embed[:,:,0,:]+self.global_embed[:,:,0:1]; x_g=x_g.permute(2,0,1); features.extend([x,x_g])
         if self.with_right_left_sensors:
-            # Left view processing
-            left_image_token, left_image_token_global = self.rgb_patch_embed(left_image)
-            if self.use_view_embed:
-                left_image_token = (
-                    left_image_token
-                    + self.view_embed[:, :, 1:2, :]
-                    + self.position_encoding(left_image_token)
-                )
-            else:
-                left_image_token = left_image_token + self.position_encoding(
-                    left_image_token
-                )
-            left_image_token = left_image_token.flatten(2).permute(2, 0, 1)
-            left_image_token_global = (
-                left_image_token_global
-                + self.view_embed[:, :, 1, :]
-                + self.global_embed[:, :, 1:2]
-            )
-            left_image_token_global = left_image_token_global.permute(2, 0, 1)
-
-            # Right view processing
-            right_image_token, right_image_token_global = self.rgb_patch_embed(
-                right_image
-            )
-            if self.use_view_embed:
-                right_image_token = (
-                    right_image_token
-                    + self.view_embed[:, :, 2:3, :]
-                    + self.position_encoding(right_image_token)
-                )
-            else:
-                right_image_token = right_image_token + self.position_encoding(
-                    right_image_token
-                )
-            right_image_token = right_image_token.flatten(2).permute(2, 0, 1)
-            right_image_token_global = (
-                right_image_token_global
-                + self.view_embed[:, :, 2, :]
-                + self.global_embed[:, :, 2:3]
-            )
-            right_image_token_global = right_image_token_global.permute(2, 0, 1)
-
-            features.extend(
-                [
-                    left_image_token,
-                    left_image_token_global,
-                    right_image_token,
-                    right_image_token_global,
-                ]
-            )
-
-        if self.with_center_sensor:
-            # Front center view processing
-            (
-                front_center_image_token,
-                front_center_image_token_global,
-            ) = self.rgb_patch_embed(front_center_image)
-            if self.use_view_embed:
-                front_center_image_token = (
-                    front_center_image_token
-                    + self.view_embed[:, :, 3:4, :]
-                    + self.position_encoding(front_center_image_token)
-                )
-            else:
-                front_center_image_token = (
-                    front_center_image_token
-                    + self.position_encoding(front_center_image_token)
-                )
-
-            front_center_image_token = front_center_image_token.flatten(2).permute(
-                2, 0, 1
-            )
-            front_center_image_token_global = (
-                front_center_image_token_global
-                + self.view_embed[:, :, 3, :]
-                + self.global_embed[:, :, 3:4]
-            )
-            front_center_image_token_global = front_center_image_token_global.permute(
-                2, 0, 1
-            )
-            features.extend([front_center_image_token, front_center_image_token_global])
-
+            x, x_g = self.rgb_patch_embed(left_image); x = x + self.view_embed[:,:,1:2,:] + self.position_encoding(x); x=x.flatten(2).permute(2,0,1); x_g=x_g+self.view_embed[:,:,1,:]+self.global_embed[:,:,1:2]; x_g=x_g.permute(2,0,1); features.extend([x,x_g])
+            x, x_g = self.rgb_patch_embed(right_image); x = x + self.view_embed[:,:,2:3,:] + self.position_encoding(x); x=x.flatten(2).permute(2,0,1); x_g=x_g+self.view_embed[:,:,2,:]+self.global_embed[:,:,2:3]; x_g=x_g.permute(2,0,1); features.extend([x,x_g])
         if self.with_lidar:
-            lidar_token, lidar_token_global = self.lidar_patch_embed(lidar)
-            if self.use_view_embed:
-                lidar_token = (
-                    lidar_token
-                    + self.view_embed[:, :, 4:5, :]
-                    + self.position_encoding(lidar_token)
-                )
-            else:
-                lidar_token = lidar_token + self.position_encoding(lidar_token)
-            lidar_token = lidar_token.flatten(2).permute(2, 0, 1)
-            lidar_token_global = (
-                lidar_token_global
-                + self.view_embed[:, :, 4, :]
-                + self.global_embed[:, :, 4:5]
-            )
-            lidar_token_global = lidar_token_global.permute(2, 0, 1)
-            features.extend([lidar_token, lidar_token_global])
-
-        features = torch.cat(features, 0)
-        return features
-
-    def forward(self, x):
-        front_image = x["rgb"]
-        left_image = x["rgb_left"]
-        right_image = x["rgb_right"]
-        front_center_image = x["rgb_center"]
-        measurements = x["measurements"]
-        target_point = x["target_point"]
-        lidar = x["lidar"]
-
+            x, x_g = self.lidar_patch_embed(lidar); x = x + self.view_embed[:,:,4:5,:] + self.position_encoding(x); x=x.flatten(2).permute(2,0,1); x_g=x_g+self.view_embed[:,:,4,:]+self.global_embed[:,:,4:5]; x_g=x_g.permute(2,0,1); features.extend([x,x_g])
+        return torch.cat(features, 0)
+    
+    def forward(self, rgb, rgb_left, rgb_right, rgb_center, lidar, measurements, target_point, **kwargs):
+        front_image=rgb
         if self.direct_concat:
-            img_size = front_image.shape[-1]
-            left_image = torch.nn.functional.interpolate(
-                left_image, size=(img_size, img_size)
-            )
-            right_image = torch.nn.functional.interpolate(
-                right_image, size=(img_size, img_size)
-            )
-            front_center_image = torch.nn.functional.interpolate(
-                front_center_image, size=(img_size, img_size)
-            )
-            front_image = torch.cat(
-                [front_image, left_image, right_image, front_center_image], dim=1
-            )
-        features = self.forward_features(
-            front_image,
-            left_image,
-            right_image,
-            front_center_image,
-            lidar,
-            measurements,
-        )
-
+            img_size=front_image.shape[-1]; front_image=torch.cat([front_image,F.interpolate(rgb_left,s:=(img_size,img_size)),F.interpolate(rgb_right,s),F.interpolate(rgb_center,s)],dim=1)
+        
+        features = self.forward_features(front_image, rgb_left, rgb_right, lidar)
         bs = front_image.shape[0]
-
-        if self.end2end:
-            tgt = self.query_pos_embed.repeat(bs, 1, 1)
-        else:
-            tgt = self.position_encoding(
-                torch.ones((bs, 1, 20, 20), device=x["rgb"].device)
-            )
-            tgt = tgt.flatten(2)
-            tgt = torch.cat([tgt, self.query_pos_embed.repeat(bs, 1, 1)], 2)
-        tgt = tgt.permute(2, 0, 1)
-
-        memory = self.encoder(features, mask=self.attn_mask)
-        hs = self.decoder(self.query_embed.repeat(1, bs, 1), memory, query_pos=tgt)[0]
-
-        hs = hs.permute(1, 0, 2)  # Batchsize ,  N, C
-        if self.end2end:
-            waypoints = self.waypoints_generator(hs, target_point)
-            return waypoints
-
-        if self.waypoints_pred_head != "heatmap":
-            traffic_feature = hs[:, :400]
-            is_junction_feature = hs[:, 400]
-            traffic_light_state_feature = hs[:, 400]
-            stop_sign_feature = hs[:, 400]
-            waypoints_feature = hs[:, 401:411]
-        else:
-            traffic_feature = hs[:, :400]
-            is_junction_feature = hs[:, 400]
-            traffic_light_state_feature = hs[:, 400]
-            stop_sign_feature = hs[:, 400]
-            waypoints_feature = hs[:, 401:405]
-
-        if self.waypoints_pred_head == "heatmap":
-            waypoints = self.waypoints_generator(waypoints_feature, measurements)
-        elif self.waypoints_pred_head == "gru":
-            waypoints = self.waypoints_generator(waypoints_feature, target_point)
-        elif self.waypoints_pred_head == "gru-command":
-            waypoints = self.waypoints_generator(waypoints_feature, target_point, measurements)
-        elif self.waypoints_pred_head == "linear":
-            waypoints = self.waypoints_generator(waypoints_feature, measurements)
-        elif self.waypoints_pred_head == "linear-sum":
-            waypoints = self.waypoints_generator(waypoints_feature, measurements)
-
+        tgt = self.position_encoding(torch.ones((bs, 1, 20, 20), device=rgb.device)).flatten(2)
+        tgt = torch.cat([tgt, self.query_pos_embed.repeat(bs, 1, 1)], 2).permute(2, 0, 1)
+        hs = self.decoder(self.query_embed.repeat(1, bs, 1), self.encoder(features), query_pos=tgt)[0].permute(1, 0, 2)
+        
+        traffic_feature = hs[:, :400]; waypoints_feature = hs[:, 401:411]; is_junction_feature = hs[:, 400]
+        waypoints = self.waypoints_generator(waypoints_feature, target_point)
         is_junction = self.junction_pred_head(is_junction_feature)
-        traffic_light_state = self.traffic_light_pred_head(traffic_light_state_feature)
-        stop_sign = self.stop_sign_head(stop_sign_feature)
-
-        velocity = measurements[:, 6:7].unsqueeze(-1)
-        velocity = velocity.repeat(1, 400, 32)
-        traffic_feature_with_vel = torch.cat([traffic_feature, velocity], dim=2)
-        traffic = self.traffic_pred_head(traffic_feature_with_vel)
+        traffic_light_state = self.traffic_light_pred_head(is_junction_feature)
+        stop_sign = self.stop_sign_head(is_junction_feature)
+        traffic = self.traffic_pred_head(traffic_feature)
+        
         return traffic, waypoints, is_junction, traffic_light_state, stop_sign, traffic_feature