Final fix: Upload self-contained model with corrected 12-channel config and full code

config.json

@@ -13,5 +13,8 @@
   "rgb_backbone_name": "r50",
   "lidar_backbone_name": "r18",
   "use_different_backbone": true,
-  "waypoints_pred_head": "gru"
+  "waypoints_pred_head": "gru",
+  "auto_map": {
+    "AutoModel": "modeling_interfuser.InterfuserForHuggingFace"
+  }
 }

modeling_interfuser.py

@@ -1,85 +1,122 @@
 
 # -*- coding: utf-8 -*-
+# This is a self-contained file for the Interfuser model.
 import torch, math, copy
 from torch import nn, Tensor
 from functools import partial
 from typing import Optional, List
+from collections import OrderedDict
 from transformers import PreTrainedModel, PretrainedConfig
-try: from timm.models.layers import to_2tuple
-except ImportError: raise ImportError("Please install timm 'pip install timm==0.3.2'")
-try: from timm.models.resnet import resnet50d, resnet26d, resnet18d
-except ImportError: raise ImportError("Please install timm 'pip install timm==0.3.2'")
+
+try:
+    from timm.models.layers import to_2tuple
+    from timm.models.resnet import resnet50d, resnet26d, resnet18d
+except ImportError:
+    raise ImportError("This model requires timm. Please install with 'pip install timm==0.3.2'")
+
+# ==================================================
+# SECTION 1: ALL DEPENDENCY CLASSES
+# ==================================================
+def _get_clones(module, N): return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
 
 class HybridEmbed(nn.Module):
     def __init__(self, backbone, img_size=224, in_chans=3, embed_dim=768, **kwargs):
         super().__init__()
-        self.backbone=backbone
+        self.backbone = backbone
         with torch.no_grad():
-            dummy_input=torch.zeros(1,in_chans,img_size,img_size)
-            features=self.backbone(dummy_input)
-            feature_dim=features[-1].shape[1]
-        self.proj=nn.Conv2d(feature_dim,embed_dim,kernel_size=1)
-    def forward(self,x):
-        features=self.backbone(x);x_feat=features[-1];x=self.proj(x_feat);global_x=torch.mean(x,[2,3],keepdim=False)[:,:,None]
-        return x,global_x
+            dummy_input = torch.zeros(1, in_chans, img_size, img_size)
+            features = self.backbone(dummy_input)
+            feature_dim = features[-1].shape[1]
+        self.proj = nn.Conv2d(feature_dim, embed_dim, kernel_size=1)
+    def forward(self, x):
+        features = self.backbone(x); x_feat = features[-1]
+        x = self.proj(x_feat); global_x = torch.mean(x, [2, 3], keepdim=False)[:, :, None]
+        return x, global_x
+
 class PositionEmbeddingSine(nn.Module):
-    def __init__(self,num_pos_feats=64,temperature=10000,normalize=True,**kwargs):
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, **kwargs):
         super().__init__();self.num_pos_feats=num_pos_feats;self.temperature=temperature;self.normalize=normalize;self.scale=2*math.pi
-    def forward(self,tensor):
+    def forward(self, tensor):
         x=tensor;bs,_,h,w=x.shape;not_mask=torch.ones((bs,h,w),device=x.device);y_embed=not_mask.cumsum(1,dtype=torch.float32);x_embed=not_mask.cumsum(2,dtype=torch.float32)
         if self.normalize:eps=1e-6;y_embed=y_embed/(y_embed[:,-1:,:]+eps)*self.scale;x_embed=x_embed/(x_embed[:,:,-1:]+eps)*self.scale
         dim_t=torch.arange(self.num_pos_feats,dtype=torch.float32,device=x.device);dim_t=self.temperature**(2*(dim_t//2)/self.num_pos_feats)
         pos_x=x_embed[:,:,:,None]/dim_t;pos_y=y_embed[:,:,:,None]/dim_t
         pos_x=torch.stack((pos_x[:,:,:,0::2].sin(),pos_x[:,:,:,1::2].cos()),dim=4).flatten(3);pos_y=torch.stack((pos_y[:,:,:,0::2].sin(),pos_y[:,:,:,1::2].cos()),dim=4).flatten(3)
         return torch.cat((pos_y,pos_x),dim=3).permute(0,3,1,2)
-def _get_clones(module,N): return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
+
 class TransformerEncoderLayer(nn.Module):
     def __init__(self,d_model,nhead,dim_feedforward=2048,dropout=0.1,activation=nn.ReLU(),**kwargs):
-        super().__init__();self.self_attn=nn.MultiheadAttention(d_model,nhead,dropout=dropout);self.linear1=nn.Linear(d_model,dim_feedforward);self.dropout=nn.Dropout(dropout);self.linear2=nn.Linear(dim_feedforward,d_model);self.norm1=nn.LayerNorm(d_model);self.norm2=nn.LayerNorm(d_model);self.dropout1=nn.Dropout(dropout);self.dropout2=nn.Dropout(dropout);self.activation=activation
+        super().__init__();self.self_attn=nn.MultiheadAttention(d_model,nhead,dropout=dropout);self.linear1=nn.Linear(d_model,dim_feedforward);self.dropout=nn.Dropout(dropout);self.linear2=nn.Linear(dim_feedforward,d_model);self.norm1=nn.LayerNorm(d_model);self.norm2=nn.LayerNorm(d_model);self.dropout1=nn.Dropout(dropout);self.dropout2=nn.Dropout(dropout);self.activation=activation()
     def with_pos_embed(self,tensor,pos:Optional[Tensor]):return tensor if pos is None else tensor+pos
     def forward(self,src,src_mask:Optional[Tensor]=None,src_key_padding_mask:Optional[Tensor]=None,pos:Optional[Tensor]=None):
         q=k=self.with_pos_embed(src,pos);src2=self.self_attn(q,k,value=src,attn_mask=src_mask,key_padding_mask=src_key_padding_mask)[0];src=src+self.dropout1(src2);src=self.norm1(src);src2=self.linear2(self.dropout(self.activation(self.linear1(src))));src=src+self.dropout2(src2);src=self.norm2(src);return src
+
 class TransformerEncoder(nn.Module):
     def __init__(self,encoder_layer,num_layers,norm=None):super().__init__();self.layers=_get_clones(encoder_layer,num_layers);self.num_layers=num_layers;self.norm=norm
     def forward(self,src,mask:Optional[Tensor]=None,src_key_padding_mask:Optional[Tensor]=None,pos:Optional[Tensor]=None):
-        output=src;[output:=layer(output,src_mask=mask,src_key_padding_mask=src_key_padding_mask,pos=pos) for layer in self.layers];
+        output=src
+        for layer in self.layers: output = layer(output,src_mask=mask,src_key_padding_mask=src_key_padding_mask,pos=pos)
         if self.norm is not None:output=self.norm(output)
         return output
+
 class TransformerDecoderLayer(nn.Module):
     def __init__(self,d_model,nhead,dim_feedforward=2048,dropout=0.1,activation=nn.ReLU(),**kwargs):
-        super().__init__();self.self_attn=nn.MultiheadAttention(d_model,nhead,dropout=dropout);self.multihead_attn=nn.MultiheadAttention(d_model,nhead,dropout=dropout);self.linear1=nn.Linear(d_model,dim_feedforward);self.dropout=nn.Dropout(dropout);self.linear2=nn.Linear(dim_feedforward,d_model);self.norm1=nn.LayerNorm(d_model);self.norm2=nn.LayerNorm(d_model);self.norm3=nn.LayerNorm(d_model);self.dropout1=nn.Dropout(dropout);self.dropout2=nn.Dropout(dropout);self.dropout3=nn.Dropout(dropout);self.activation=activation
+        super().__init__();self.self_attn=nn.MultiheadAttention(d_model,nhead,dropout=dropout);self.multihead_attn=nn.MultiheadAttention(d_model,nhead,dropout=dropout);self.linear1=nn.Linear(d_model,dim_feedforward);self.dropout=nn.Dropout(dropout);self.linear2=nn.Linear(dim_feedforward,d_model);self.norm1=nn.LayerNorm(d_model);self.norm2=nn.LayerNorm(d_model);self.norm3=nn.LayerNorm(d_model);self.dropout1=nn.Dropout(dropout);self.dropout2=nn.Dropout(dropout);self.dropout3=nn.Dropout(dropout);self.activation=activation()
     def with_pos_embed(self,tensor,pos:Optional[Tensor]):return tensor if pos is None else tensor+pos
     def forward(self,tgt,memory,tgt_mask:Optional[Tensor]=None,memory_mask:Optional[Tensor]=None,tgt_key_padding_mask:Optional[Tensor]=None,memory_key_padding_mask:Optional[Tensor]=None,pos:Optional[Tensor]=None,query_pos:Optional[Tensor]=None):
-        q=k=self.with_pos_embed(tgt,query_pos);tgt2=self.self_attn(q,k,value=tgt,attn_mask=tgt_mask,key_padding_mask=tgt_key_padding_mask)[0];tgt=tgt+self.dropout1(tgt2);tgt=self.norm1(tgt);tgt2=self.multihead_attn(query=self.with_pos_embed(tgt,query_pos),key=self.with_pos_embed(memory,pos),value=memory,attn_mask=memory_mask,key_padding_mask=memory_key_padding_mask)[0];tgt=tgt+self.dropout2(tgt2);tgt=self.norm2(tgt);tgt2=self.linear2(self.dropout(self.activation(self.linear1(tgt))));tgt=tgt+self.dropout3(tgt2);tgt=self.norm3(tgt);return tgt
+        q=k=self.with_pos_embed(tgt,query_pos);tgt2=self.self_attn(q,k,value=tgt,attn_mask=tgt_mask,key_padding_mask=tgt_key_padding_mask)[0];tgt=tgt+self.dropout1(tgt2);tgt=self.norm1(tgt)
+        tgt2=self.multihead_attn(query=self.with_pos_embed(tgt,query_pos),key=self.with_pos_embed(memory,pos),value=memory,attn_mask=memory_mask,key_padding_mask=memory_key_padding_mask)[0];tgt=tgt+self.dropout2(tgt2);tgt=self.norm2(tgt)
+        tgt2=self.linear2(self.dropout(self.activation(self.linear1(tgt))));tgt=tgt+self.dropout3(tgt2);tgt=self.norm3(tgt);return tgt
+
 class TransformerDecoder(nn.Module):
     def __init__(self,decoder_layer,num_layers,norm=None,return_intermediate=False):super().__init__();self.layers=_get_clones(decoder_layer,num_layers);self.num_layers=num_layers;self.norm=norm;self.return_intermediate=return_intermediate
-    def forward(self,tgt,memory,**kwargs):output=tgt;[output:=layer(output,memory,**kwargs) for layer in self.layers];return self.norm(output).unsqueeze(0) if self.norm is not None else output.unsqueeze(0)
+    def forward(self,tgt,memory,**kwargs):
+        output=tgt
+        for layer in self.layers: output = layer(output,memory,**kwargs)
+        if self.norm is not None:output=self.norm(output)
+        return output.unsqueeze(0)
+
 class GRUWaypointsPredictor(nn.Module):
-    def __init__(self,input_dim,waypoints=10):super().__init__();self.gru=torch.nn.GRU(input_size=input_dim,hidden_size=64,batch_first=True);self.encoder=nn.Linear(2,64);self.decoder=nn.Linear(64,2);self.waypoints=waypoints
+    def __init__(self,input_dim,waypoints=10,**kwargs):super().__init__();self.gru=torch.nn.GRU(input_size=input_dim,hidden_size=64,batch_first=True);self.encoder=nn.Linear(2,64);self.decoder=nn.Linear(64,2);self.waypoints=waypoints
     def forward(self,x,target_point):bs=x.shape[0];z=self.encoder(target_point).unsqueeze(0);output,_=self.gru(x,z);output=self.decoder(output.reshape(bs*self.waypoints,-1)).reshape(bs,self.waypoints,2);return torch.cumsum(output,1)
 
+# ==================================================
+# SECTION 2: ORIGINAL INTERFUSER MODEL
+# ==================================================
 class Interfuser(nn.Module):
-    def __init__(self,config):
-        super().__init__();self.config=config;self.embed_dim=config.embed_dim;self.waypoints_pred_head=config.waypoints_pred_head
-        if config.rgb_backbone_name=="r50":self.rgb_backbone=resnet50d(pretrained=False,in_chans=config.in_chans,features_only=True,out_indices=[4])
-        else:self.rgb_backbone=resnet18d(pretrained=False,in_chans=config.in_chans,features_only=True,out_indices=[4])
-        self.rgb_patch_embed=partial(HybridEmbed,backbone=self.rgb_backbone)(in_chans=config.in_chans,embed_dim=config.embed_dim)
-        if getattr(config,'lidar_backbone_name',None):self.lidar_backbone=resnet18d(pretrained=False,in_chans=3,features_only=True,out_indices=[4]);self.lidar_patch_embed=partial(HybridEmbed,backbone=self.lidar_backbone)(in_chans=3,embed_dim=config.embed_dim)
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.traffic_pred_head_type=getattr(config,"traffic_pred_head_type","det");self.embed_dim=config.embed_dim;self.waypoints_pred_head=config.waypoints_pred_head;self.use_view_embed=getattr(config,"use_view_embed",True)
+        if config.use_different_backbone:
+            if config.rgb_backbone_name=="r50":self.rgb_backbone=resnet50d(pretrained=False,in_chans=config.in_chans,features_only=True,out_indices=[4])
+            elif config.rgb_backbone_name=="r26":self.rgb_backbone=resnet26d(pretrained=False,in_chans=config.in_chans,features_only=True,out_indices=[4])
+            else:self.rgb_backbone=resnet18d(pretrained=False,in_chans=config.in_chans,features_only=True,out_indices=[4])
+            if config.lidar_backbone_name=="r18":self.lidar_backbone=resnet18d(pretrained=False,in_chans=3,features_only=True,out_indices=[4])
+            self.rgb_patch_embed=partial(HybridEmbed,backbone=self.rgb_backbone)(in_chans=config.in_chans,embed_dim=config.embed_dim)
+            self.lidar_patch_embed=partial(HybridEmbed,backbone=self.lidar_backbone)(in_chans=3,embed_dim=config.embed_dim)
         self.global_embed=nn.Parameter(torch.zeros(1,config.embed_dim,5));self.view_embed=nn.Parameter(torch.zeros(1,config.embed_dim,5,1));self.query_pos_embed=nn.Parameter(torch.zeros(1,config.embed_dim,11));self.query_embed=nn.Parameter(torch.zeros(400+11,1,config.embed_dim))
         self.waypoints_generator=GRUWaypointsPredictor(config.embed_dim);self.junction_pred_head=nn.Linear(config.embed_dim,2);self.traffic_light_pred_head=nn.Linear(config.embed_dim,2);self.stop_sign_head=nn.Linear(config.embed_dim,2)
         self.traffic_pred_head=nn.Sequential(nn.Linear(config.embed_dim+32,64),nn.ReLU(),nn.Linear(64,7),nn.Sigmoid());self.position_encoding=PositionEmbeddingSine(config.embed_dim//2)
-        encoder_layer=TransformerEncoderLayer(config.embed_dim,config.num_heads,config.dim_feedforward);self.encoder=TransformerEncoder(encoder_layer,config.enc_depth);decoder_layer=TransformerDecoderLayer(config.embed_dim,config.num_heads,config.dim_feedforward);self.decoder=TransformerDecoder(decoder_layer,config.dec_depth,nn.LayerNorm(config.embed_dim))
+        encoder_layer=TransformerEncoderLayer(config.embed_dim,config.num_heads,config.dim_feedforward);self.encoder=TransformerEncoder(encoder_layer,config.enc_depth)
+        decoder_layer=TransformerDecoderLayer(config.embed_dim,config.num_heads,config.dim_feedforward);self.decoder=TransformerDecoder(decoder_layer,config.dec_depth,nn.LayerNorm(config.embed_dim))
     def forward_features(self,front_image,left_image,right_image,front_center_image,lidar,**kwargs):
-        features=[];front_image_token,_=self.rgb_patch_embed(front_image);features.append(front_image_token+self.position_encoding(front_image_token))
-        left_image_token,_=self.rgb_patch_embed(left_image);features.append(left_image_token+self.position_encoding(left_image_token))
-        right_image_token,_=self.rgb_patch_embed(right_image);features.append(right_image_token+self.position_encoding(right_image_token))
-        return torch.cat([f.flatten(2) for f in features],-1).permute(2,0,1)
+        features=[];front_image_token,front_image_token_global=self.rgb_patch_embed(front_image)
+        front_image_token=(front_image_token+self.view_embed[:,:,0:1,:]+self.position_encoding(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]).permute(2,0,1)
+        features.extend([front_image_token,front_image_token_global])
+        left_image_token,left_image_token_global=self.rgb_patch_embed(left_image)
+        left_image_token=(left_image_token+self.view_embed[:,:,1:2,:]+self.position_encoding(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]).permute(2,0,1)
+        features.extend([left_image_token,left_image_token_global])
+        right_image_token,right_image_token_global=self.rgb_patch_embed(right_image)
+        right_image_token=(right_image_token+self.view_embed[:,:,2:3,:]+self.position_encoding(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]).permute(2,0,1)
+        features.extend([right_image_token,right_image_token_global])
+        return torch.cat(features,0)
     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"]
         if self.config.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))
-            input_image=torch.cat([front_image,left_image,right_image,front_center_image],dim=1)
-            features,_=self.rgb_patch_embed(input_image);features=features+self.position_encoding(features);features=features.flatten(2).permute(2,0,1)
+            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=measurements)
         bs=front_image.shape[0];tgt=self.position_encoding(torch.ones((bs,1,20,20),device=x["rgb"].device)).flatten(2);tgt=torch.cat([tgt,self.query_pos_embed.repeat(bs,1,1)],2).permute(2,0,1)
         memory=self.encoder(features,mask=None);hs=self.decoder(self.query_embed.repeat(1,bs,1),memory,query_pos=tgt)[0].permute(1,0,2)
         waypoints=self.waypoints_generator(hs[:,401:411],target_point)
@@ -87,17 +124,23 @@ class Interfuser(nn.Module):
         velocity=measurements[:,6:7].unsqueeze(-1).repeat(1,400,32);traffic_feature_with_vel=torch.cat([hs[:,:400],velocity],dim=2);traffic=self.traffic_pred_head(traffic_feature_with_vel)
         return traffic,waypoints,is_junction,traffic_light_state,stop_sign,hs[:,:400]
 
+# ==================================================
+# SECTION 3: HUGGING FACE WRAPPER CLASSES
+# ==================================================
 class InterfuserConfig(PretrainedConfig):
     model_type="interfuser"
-    def __init__(self,in_chans=12,embed_dim=256,enc_depth=6,dec_depth=6,num_heads=8,dim_feedforward=2048,rgb_backbone_name="r50",lidar_backbone_name="r18",use_different_backbone=True,waypoints_pred_head="gru",direct_concat=True,**kwargs):
-        super().__init__(**kwargs);self.in_chans=in_chans;self.embed_dim,self.enc_depth,self.dec_depth,self.num_heads,self.dim_feedforward=embed_dim,enc_depth,dec_depth,num_heads,dim_feedforward
-        self.rgb_backbone_name,self.lidar_backbone_name,self.use_different_backbone,self.waypoints_pred_head,self.direct_concat=rgb_backbone_name,lidar_backbone_name,use_different_backbone,waypoints_pred_head,direct_concat
-        self.auto_map={"AutoModel":"modeling_interfuser.InterfuserForHuggingFace"}
+    def __init__(self, in_chans=12, embed_dim=256, enc_depth=6, dec_depth=6, num_heads=8, dim_feedforward=2048, rgb_backbone_name="r50", lidar_backbone_name="r18", use_different_backbone=True, waypoints_pred_head="gru", direct_concat=True, **kwargs):
+        super().__init__(**kwargs)
+        self.in_chans = in_chans; self.embed_dim, self.enc_depth, self.dec_depth, self.num_heads, self.dim_feedforward = embed_dim, enc_depth, dec_depth, num_heads, dim_feedforward
+        self.rgb_backbone_name, self.lidar_backbone_name, self.use_different_backbone, self.waypoints_pred_head, self.direct_concat = rgb_backbone_name, lidar_backbone_name, use_different_backbone, waypoints_pred_head, direct_concat
+        self.auto_map = {"AutoModel": "modeling_interfuser.InterfuserForHuggingFace"}
 
 class InterfuserForHuggingFace(PreTrainedModel):
-    config_class=InterfuserConfig
-    def __init__(self,config:InterfuserConfig):
-        super().__init__(config);self.config=config;self.interfuser_model=Interfuser(config)
-    def forward(self,rgb,rgb_left,rgb_right,rgb_center,lidar,measurements,target_point,**kwargs):
-        inputs_dict={'rgb':rgb,'rgb_left':rgb_left,'rgb_right':rgb_right,'rgb_center':rgb_center,'lidar':lidar,'measurements':measurements,'target_point':target_point}
+    config_class = InterfuserConfig
+    def __init__(self, config: InterfuserConfig):
+        super().__init__(config)
+        self.config = config
+        self.interfuser_model = Interfuser(config)
+    def forward(self, rgb, rgb_left, rgb_right, rgb_center, lidar, measurements, target_point, **kwargs):
+        inputs_dict = {'rgb':rgb, 'rgb_left':rgb_left, 'rgb_right':right_image, 'rgb_center':rgb_center, 'lidar':lidar, 'measurements':measurements, 'target_point':target_point}
         return self.interfuser_model.forward(inputs_dict)