Update references

Author: jainapurva

benchmarks/_models/llama/eval.py

@@ -15,7 +15,7 @@
 from tokenizer import get_tokenizer
 
 import torchao
-from torchao._models.llama.model import prepare_inputs_for_model
+from benchmarks._models.llama.model import prepare_inputs_for_model
 from torchao.quantization import (
     PerRow,
     PerTensor,
@@ -172,7 +172,7 @@ def run_evaluation(
         if "autoround" in quantization:
             from transformers import AutoTokenizer
 
-            from torchao._models.llama.model import TransformerBlock
+            from benchmarks._models.llama.model import TransformerBlock
             from torchao.prototype.autoround.autoround_llm import (
                 quantize_model_with_autoround_,
             )

benchmarks/_models/llama/generate.py

@@ -14,7 +14,7 @@
 import torch._inductor.config
 
 import torchao
-from torchao._models.utils import (
+from benchmarks._models.utils import (
     get_arch_name,
     write_json_result_local,
     write_json_result_ossci,
@@ -72,8 +72,8 @@ def device_sync(device):
 wd = Path(__file__).parent.parent.resolve()
 sys.path.append(str(wd))
 
-from torchao._models.llama.model import Transformer, prepare_inputs_for_model
-from torchao._models.llama.tokenizer import get_tokenizer
+from benchmarks._models.llama.model import Transformer, prepare_inputs_for_model
+from benchmarks._models.llama.tokenizer import get_tokenizer
 
 
 def multinomial_sample_one_no_sync(
@@ -575,8 +575,8 @@ def ffn_or_attn_only(mod, fqn):
                 model, float8_dynamic_activation_float8_weight(granularity=granularity)
             )
         elif "autoquant_v2" in quantization:
+            from benchmarks._models.llama.model import prepare_inputs_for_model
             from torchao._eval import InputRecorder
-            from torchao._models.llama.model import prepare_inputs_for_model
             from torchao.prototype.quantization.autoquant_v2 import autoquant_v2
 
             calibration_seq_length = 256
@@ -665,8 +665,8 @@ def ffn_or_attn_only(mod, fqn):
             # do autoquantization
             model.finalize_autoquant()
         elif "autoquant" in quantization:
+            from benchmarks._models.llama.model import prepare_inputs_for_model
             from torchao._eval import InputRecorder
-            from torchao._models.llama.model import prepare_inputs_for_model
 
             calibration_seq_length = 256
             inputs = (

benchmarks/_models/llama/perf_profile.py

@@ -116,8 +116,8 @@
 import torch
 from torch.nn.attention import SDPBackend
 
-from torchao._models.llama.model import Transformer
-from torchao._models.llama.tokenizer import get_tokenizer
+from benchmarks._models.llama.model import Transformer
+from benchmarks._models.llama.tokenizer import get_tokenizer
 from torchao.prototype.profiler import (
     CUDADeviceSpec,
     TransformerPerformanceCounter,

benchmarks/_models/sam/eval_combo.py

@@ -9,7 +9,7 @@
 from metrics import calculate_miou, create_result_entry
 
 import torchao
-from torchao._models.utils import (
+from benchmarks._models.utils import (
     get_arch_name,
     write_json_result_local,
     write_json_result_ossci,

benchmarks/_models/sam2/__init__.py

@@ -8,4 +8,4 @@
 from hydra.core.global_hydra import GlobalHydra
 
 if not GlobalHydra.instance().is_initialized():
-    initialize_config_module("torchao._models.sam2", version_base="1.2")
+    initialize_config_module("benchmarks._models.sam2", version_base="1.2")

benchmarks/_models/sam2/automatic_mask_generator.py

@@ -11,9 +11,9 @@
 import torch
 from torchvision.ops.boxes import batched_nms, box_area  # type: ignore
 
-from torchao._models.sam2.modeling.sam2_base import SAM2Base
-from torchao._models.sam2.sam2_image_predictor import SAM2ImagePredictor
-from torchao._models.sam2.utils.amg import (
+from benchmarks._models.sam2.modeling.sam2_base import SAM2Base
+from benchmarks._models.sam2.sam2_image_predictor import SAM2ImagePredictor
+from benchmarks._models.sam2.utils.amg import (
     MaskData,
     _mask_to_rle_pytorch_2_0,
     _mask_to_rle_pytorch_2_1,
@@ -33,7 +33,7 @@
     uncrop_masks,
     uncrop_points,
 )
-from torchao._models.sam2.utils.misc import (
+from benchmarks._models.sam2.utils.misc import (
     crop_image,
     get_image_size,
 )

benchmarks/_models/sam2/build_sam.py

@@ -12,7 +12,7 @@
 from hydra.utils import instantiate
 from omegaconf import OmegaConf
 
-from torchao._models import sam2
+from benchmarks._models import sam2
 
 # Check if the user is running Python from the parent directory of the sam2 repo
 # (i.e. the directory where this repo is cloned into) -- this is not supported since
@@ -106,7 +106,7 @@ def build_sam2_video_predictor(
     **kwargs,
 ):
     hydra_overrides = [
-        "++model._target_=torchao._models.sam2.sam2_video_predictor.SAM2VideoPredictor",
+        "++model._target_=benchmarks._models.sam2.sam2_video_predictor.SAM2VideoPredictor",
     ]
     if apply_postprocessing:
         hydra_overrides_extra = hydra_overrides_extra.copy()

benchmarks/_models/sam2/configs/sam2.1/sam2.1_hiera_b+.yaml

@@ -2,18 +2,18 @@
 
 # Model
 model:
-  _target_: torchao._models.sam2.modeling.sam2_base.SAM2Base
+  _target_: benchmarks._models.sam2.modeling.sam2_base.SAM2Base
   image_encoder:
-    _target_: torchao._models.sam2.modeling.backbones.image_encoder.ImageEncoder
+    _target_: benchmarks._models.sam2.modeling.backbones.image_encoder.ImageEncoder
     scalp: 1
     trunk:
-      _target_: torchao._models.sam2.modeling.backbones.hieradet.Hiera
+      _target_: benchmarks._models.sam2.modeling.backbones.hieradet.Hiera
       embed_dim: 112
       num_heads: 2
     neck:
-      _target_: torchao._models.sam2.modeling.backbones.image_encoder.FpnNeck
+      _target_: benchmarks._models.sam2.modeling.backbones.image_encoder.FpnNeck
       position_encoding:
-        _target_: torchao._models.sam2.modeling.position_encoding.PositionEmbeddingSine
+        _target_: benchmarks._models.sam2.modeling.position_encoding.PositionEmbeddingSine
         num_pos_feats: 256
         normalize: true
         scale: null
@@ -24,17 +24,17 @@ model:
       fpn_interp_model: nearest
 
   memory_attention:
-    _target_: torchao._models.sam2.modeling.memory_attention.MemoryAttention
+    _target_: benchmarks._models.sam2.modeling.memory_attention.MemoryAttention
     d_model: 256
     pos_enc_at_input: true
     layer:
-      _target_: torchao._models.sam2.modeling.memory_attention.MemoryAttentionLayer
+      _target_: benchmarks._models.sam2.modeling.memory_attention.MemoryAttentionLayer
       activation: relu
       dim_feedforward: 2048
       dropout: 0.1
       pos_enc_at_attn: false
       self_attention:
-        _target_: torchao._models.sam2.modeling.sam.transformer.RoPEAttention
+        _target_: benchmarks._models.sam2.modeling.sam.transformer.RoPEAttention
         rope_theta: 10000.0
         feat_sizes: [32, 32]
         embedding_dim: 256
@@ -45,7 +45,7 @@ model:
       pos_enc_at_cross_attn_keys: true
       pos_enc_at_cross_attn_queries: false
       cross_attention:
-        _target_: torchao._models.sam2.modeling.sam.transformer.RoPEAttention
+        _target_: benchmarks._models.sam2.modeling.sam.transformer.RoPEAttention
         rope_theta: 10000.0
         feat_sizes: [32, 32]
         rope_k_repeat: True
@@ -57,23 +57,23 @@ model:
     num_layers: 4
 
   memory_encoder:
-      _target_: torchao._models.sam2.modeling.memory_encoder.MemoryEncoder
+      _target_: benchmarks._models.sam2.modeling.memory_encoder.MemoryEncoder
       out_dim: 64
       position_encoding:
-        _target_: torchao._models.sam2.modeling.position_encoding.PositionEmbeddingSine
+        _target_: benchmarks._models.sam2.modeling.position_encoding.PositionEmbeddingSine
         num_pos_feats: 64
         normalize: true
         scale: null
         temperature: 10000
       mask_downsampler:
-        _target_: torchao._models.sam2.modeling.memory_encoder.MaskDownSampler
+        _target_: benchmarks._models.sam2.modeling.memory_encoder.MaskDownSampler
         kernel_size: 3
         stride: 2
         padding: 1
       fuser:
-        _target_: torchao._models.sam2.modeling.memory_encoder.Fuser
+        _target_: benchmarks._models.sam2.modeling.memory_encoder.Fuser
         layer:
-          _target_: torchao._models.sam2.modeling.memory_encoder.CXBlock
+          _target_: benchmarks._models.sam2.modeling.memory_encoder.CXBlock
           dim: 256
           kernel_size: 7
           padding: 3

benchmarks/_models/sam2/configs/sam2.1/sam2.1_hiera_l.yaml

@@ -2,22 +2,22 @@
 
 # Model
 model:
-  _target_: torchao._models.sam2.modeling.sam2_base.SAM2Base
+  _target_: benchmarks._models.sam2.modeling.sam2_base.SAM2Base
   image_encoder:
-    _target_: torchao._models.sam2.modeling.backbones.image_encoder.ImageEncoder
+    _target_: benchmarks._models.sam2.modeling.backbones.image_encoder.ImageEncoder
     scalp: 1
     trunk:
-      _target_: torchao._models.sam2.modeling.backbones.hieradet.Hiera
+      _target_: benchmarks._models.sam2.modeling.backbones.hieradet.Hiera
       embed_dim: 144
       num_heads: 2
       stages: [2, 6, 36, 4]
       global_att_blocks: [23, 33, 43]
       window_pos_embed_bkg_spatial_size: [7, 7]
       window_spec: [8, 4, 16, 8]
     neck:
-      _target_: torchao._models.sam2.modeling.backbones.image_encoder.FpnNeck
+      _target_: benchmarks._models.sam2.modeling.backbones.image_encoder.FpnNeck
       position_encoding:
-        _target_: torchao._models.sam2.modeling.position_encoding.PositionEmbeddingSine
+        _target_: benchmarks._models.sam2.modeling.position_encoding.PositionEmbeddingSine
         num_pos_feats: 256
         normalize: true
         scale: null
@@ -28,17 +28,17 @@ model:
       fpn_interp_model: nearest
 
   memory_attention:
-    _target_: torchao._models.sam2.modeling.memory_attention.MemoryAttention
+    _target_: benchmarks._models.sam2.modeling.memory_attention.MemoryAttention
     d_model: 256
     pos_enc_at_input: true
     layer:
-      _target_: torchao._models.sam2.modeling.memory_attention.MemoryAttentionLayer
+      _target_: benchmarks._models.sam2.modeling.memory_attention.MemoryAttentionLayer
       activation: relu
       dim_feedforward: 2048
       dropout: 0.1
       pos_enc_at_attn: false
       self_attention:
-        _target_: torchao._models.sam2.modeling.sam.transformer.RoPEAttention
+        _target_: benchmarks._models.sam2.modeling.sam.transformer.RoPEAttention
         rope_theta: 10000.0
         feat_sizes: [32, 32]
         embedding_dim: 256
@@ -49,7 +49,7 @@ model:
       pos_enc_at_cross_attn_keys: true
       pos_enc_at_cross_attn_queries: false
       cross_attention:
-        _target_: torchao._models.sam2.modeling.sam.transformer.RoPEAttention
+        _target_: benchmarks._models.sam2.modeling.sam.transformer.RoPEAttention
         rope_theta: 10000.0
         feat_sizes: [32, 32]
         rope_k_repeat: True
@@ -61,23 +61,23 @@ model:
     num_layers: 4
 
   memory_encoder:
-      _target_: torchao._models.sam2.modeling.memory_encoder.MemoryEncoder
+      _target_: benchmarks._models.sam2.modeling.memory_encoder.MemoryEncoder
       out_dim: 64
       position_encoding:
-        _target_: torchao._models.sam2.modeling.position_encoding.PositionEmbeddingSine
+        _target_: benchmarks._models.sam2.modeling.position_encoding.PositionEmbeddingSine
         num_pos_feats: 64
         normalize: true
         scale: null
         temperature: 10000
       mask_downsampler:
-        _target_: torchao._models.sam2.modeling.memory_encoder.MaskDownSampler
+        _target_: benchmarks._models.sam2.modeling.memory_encoder.MaskDownSampler
         kernel_size: 3
         stride: 2
         padding: 1
       fuser:
-        _target_: torchao._models.sam2.modeling.memory_encoder.Fuser
+        _target_: benchmarks._models.sam2.modeling.memory_encoder.Fuser
         layer:
-          _target_: torchao._models.sam2.modeling.memory_encoder.CXBlock
+          _target_: benchmarks._models.sam2.modeling.memory_encoder.CXBlock
           dim: 256
           kernel_size: 7
           padding: 3

benchmarks/_models/sam2/configs/sam2.1/sam2.1_hiera_s.yaml

@@ -2,21 +2,21 @@
 
 # Model
 model:
-  _target_: torchao._models.sam2.modeling.sam2_base.SAM2Base
+  _target_: benchmarks._models.sam2.modeling.sam2_base.SAM2Base
   image_encoder:
-    _target_: torchao._models.sam2.modeling.backbones.image_encoder.ImageEncoder
+    _target_: benchmarks._models.sam2.modeling.backbones.image_encoder.ImageEncoder
     scalp: 1
     trunk:
-      _target_: torchao._models.sam2.modeling.backbones.hieradet.Hiera
+      _target_: benchmarks._models.sam2.modeling.backbones.hieradet.Hiera
       embed_dim: 96
       num_heads: 1
       stages: [1, 2, 11, 2]
       global_att_blocks: [7, 10, 13]
       window_pos_embed_bkg_spatial_size: [7, 7]
     neck:
-      _target_: torchao._models.sam2.modeling.backbones.image_encoder.FpnNeck
+      _target_: benchmarks._models.sam2.modeling.backbones.image_encoder.FpnNeck
       position_encoding:
-        _target_: torchao._models.sam2.modeling.position_encoding.PositionEmbeddingSine
+        _target_: benchmarks._models.sam2.modeling.position_encoding.PositionEmbeddingSine
         num_pos_feats: 256
         normalize: true
         scale: null
@@ -27,17 +27,17 @@ model:
       fpn_interp_model: nearest
 
   memory_attention:
-    _target_: torchao._models.sam2.modeling.memory_attention.MemoryAttention
+    _target_: benchmarks._models.sam2.modeling.memory_attention.MemoryAttention
     d_model: 256
     pos_enc_at_input: true
     layer:
-      _target_: torchao._models.sam2.modeling.memory_attention.MemoryAttentionLayer
+      _target_: benchmarks._models.sam2.modeling.memory_attention.MemoryAttentionLayer
       activation: relu
       dim_feedforward: 2048
       dropout: 0.1
       pos_enc_at_attn: false
       self_attention:
-        _target_: torchao._models.sam2.modeling.sam.transformer.RoPEAttention
+        _target_: benchmarks._models.sam2.modeling.sam.transformer.RoPEAttention
         rope_theta: 10000.0
         feat_sizes: [32, 32]
         embedding_dim: 256
@@ -48,7 +48,7 @@ model:
       pos_enc_at_cross_attn_keys: true
       pos_enc_at_cross_attn_queries: false
       cross_attention:
-        _target_: torchao._models.sam2.modeling.sam.transformer.RoPEAttention
+        _target_: benchmarks._models.sam2.modeling.sam.transformer.RoPEAttention
         rope_theta: 10000.0
         feat_sizes: [32, 32]
         rope_k_repeat: True
@@ -60,23 +60,23 @@ model:
     num_layers: 4
 
   memory_encoder:
-      _target_: torchao._models.sam2.modeling.memory_encoder.MemoryEncoder
+      _target_: benchmarks._models.sam2.modeling.memory_encoder.MemoryEncoder
       out_dim: 64
       position_encoding:
-        _target_: torchao._models.sam2.modeling.position_encoding.PositionEmbeddingSine
+        _target_: benchmarks._models.sam2.modeling.position_encoding.PositionEmbeddingSine
         num_pos_feats: 64
         normalize: true
         scale: null
         temperature: 10000
       mask_downsampler:
-        _target_: torchao._models.sam2.modeling.memory_encoder.MaskDownSampler
+        _target_: benchmarks._models.sam2.modeling.memory_encoder.MaskDownSampler
         kernel_size: 3
         stride: 2
         padding: 1
       fuser:
-        _target_: torchao._models.sam2.modeling.memory_encoder.Fuser
+        _target_: benchmarks._models.sam2.modeling.memory_encoder.Fuser
         layer:
-          _target_: torchao._models.sam2.modeling.memory_encoder.CXBlock
+          _target_: benchmarks._models.sam2.modeling.memory_encoder.CXBlock
           dim: 256
           kernel_size: 7
           padding: 3