original_docstring: Add an output head for models predicting "box-level outputs".
docstring: Add an output head for models predicting "box-level outputs".
docstring_tokens: Add an output head for models predicting " box - level outputs " .
code: def add_roi_head(model, blob_in, dim_in, pool_stride, out_spatial_dim,
               suffix, lfb_infer_only, test_mode):
      """Add an output head for models predicting "box-level outputs"."""
  
      # (B, 2048, 16, 14, 14) -> (N, 2048, 1, 1, 1)
      roi_feat = roi_pool(model, blob_in, dim_in, out_spatial_dim, suffix)
  
      all_heads = [roi_feat]
      new_dim_in = [dim_in]
  
      if cfg.LFB.ENABLED and not lfb_infer_only:
  
          fbo_out, fbo_out_dim = lfb_helper.add_fbo_head(
              model, roi_feat, dim_in,
              num_lfb_feat=(cfg.LFB.WINDOW_SIZE
                            * cfg.AVA.LFB_MAX_NUM_FEAT_PER_STEP),
              test_mode=test_mode, suffix=suffix)
  
          all_heads.append(fbo_out)
          new_dim_in.append(fbo_out_dim)
  
      return (model.net.Concat(all_heads,
                               ['pool5', 'pool5_concat_info'],
                               axis=1)[0],
              sum(new_dim_in))
code_tokens: def add_roi_head ( model , blob_in , dim_in , pool_stride , out_spatial_dim , suffix , lfb_infer_only , test_mode ) : roi_feat = roi_pool ( model , blob_in , dim_in , out_spatial_dim , suffix ) all_heads = [ roi_feat ] new_dim_in = [ dim_in ] if cfg . LFB . ENABLED and not lfb_infer_only : fbo_out , fbo_out_dim = lfb_helper . add_fbo_head ( model , roi_feat , dim_in , num_lfb_feat = ( cfg . LFB . WINDOW_SIZE * cfg . AVA . LFB_MAX_NUM_FEAT_PER_STEP ) , test_mode = test_mode , suffix = suffix ) all_heads . append ( fbo_out ) new_dim_in . append ( fbo_out_dim ) return ( model . net . Concat ( all_heads , [ 'pool5' , 'pool5_concat_info' ] , axis = 1 ) [ 0 ] , sum ( new_dim_in ) )
short_docstring: Add an output head for models predicting "box-level outputs".
short_docstring_tokens: Add an output head for models predicting " box - level outputs " .
comment: """Add an output head for models predicting "box-level outputs"."""$SEP$# (B, 2048, 16, 14, 14) -> (N, 2048, 1, 1, 1)