#0: Added a naive batching mechanism for some models in the tests folder

tests/conftest.py

@@ -36,7 +36,7 @@ def pytest_addoption(parser):
         help="Run up to the specified iteration count and report metrics based on this iteration.",
     )
     parser.addoption("--gen_op_accuracy_tests", action="store_true")
-
+    parser.addoption("--batch_size", action="store", default=None, help="Batch size for testing")
 
 @pytest.fixture(scope="session")
 def input_var_only_native(request):
@@ -69,6 +69,9 @@ def device():
     ttnn.synchronize_device(device)
     ttnn.close_device(device)
 
+@pytest.fixture(scope="session")
+def get_batch_size(request):
+    return request.config.getoption("--batch_size")
 
 def get_dispatch_core_type():
     # Instead of conditionally returning WORKER or ETH, here we always return ETH

tests/models/MobileNetV2/test_MobileNetV2.py

@@ -5,7 +5,8 @@
 from PIL import Image
 
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
+
 
 
 class ThisTester(ModelTester):
@@ -32,12 +33,18 @@ def _load_inputs(self):
     ["eval"],
 )
 @pytest.mark.converted_end_to_end
-def test_MobileNetV2(record_property, mode):
+def test_MobileNetV2(record_property, mode, get_batch_size):
     model_name = "MobileNetV2"
     record_property("model_name", model_name)
     record_property("mode", mode)
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size)
     if mode == "eval":
         # Print the top 5 predictions
         _, indices = torch.topk(results, 5)

tests/models/albert/test_albert_masked_lm.py

@@ -3,7 +3,7 @@
 from transformers import AutoTokenizer, AlbertForMaskedLM
 import torch
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits
 
 
 class ThisTester(ModelTester):
@@ -40,15 +40,25 @@ def append_fake_loss_function(self, outputs):
         "albert/albert-xxlarge-v2",
     ],
 )
-def test_albert_masked_lm(record_property, model_name, mode):
+
+
+def test_albert_masked_lm(record_property, model_name, mode, get_batch_size):
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
-
+    results = tester.test_model(batch_size=batch_size)
+    
     if mode == "eval":
         # retrieve index of [MASK]
+
+        results.logits = process_batched_logits(results.logits, batch_size)
+        #print(results.logits.shape)
         logits = results.logits
         mask_token_index = (tester.inputs.input_ids == tester.tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
         predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)

tests/models/albert/test_albert_question_answering.py

@@ -3,7 +3,7 @@
 from transformers import AutoTokenizer, AlbertForQuestionAnswering
 import torch
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
 
 
 class ThisTester(ModelTester):
@@ -24,17 +24,22 @@ def _load_inputs(self):
 )
 @pytest.mark.converted_end_to_end
 @pytest.mark.parametrize("model_name", ["twmkn9/albert-base-v2-squad2"])
-def test_albert_question_answering(record_property, model_name, mode):
+def test_albert_question_answering(record_property, model_name, mode, get_batch_size):
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
-        answer_start_index = results.start_logits.argmax()
-        answer_end_index = results.end_logits.argmax()
-
+        answer_start_index = process_batched_logits(results.start_logits,batch_size).argmax()
+        answer_end_index = process_batched_logits(results.end_logits,batch_size).argmax()
         predict_answer_tokens = tester.inputs.input_ids[0, answer_start_index : answer_end_index + 1]
         answer = tester.tokenizer.decode(predict_answer_tokens, skip_special_tokens=True)
 

tests/models/albert/test_albert_sequence_classification.py

@@ -3,7 +3,7 @@
 from transformers import AlbertTokenizer, AlbertForSequenceClassification
 import torch
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
 
 
 class ThisTester(ModelTester):
@@ -23,15 +23,21 @@ def _load_inputs(self):
 )
 @pytest.mark.converted_end_to_end
 @pytest.mark.parametrize("model_name", ["textattack/albert-base-v2-imdb"])
-def test_albert_sequence_classification(record_property, model_name, mode):
+def test_albert_sequence_classification(record_property, model_name, mode, get_batch_size):
     record_property("model_name", model_name)
     record_property("mode", mode)
 
-    tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
 
+    tester = ThisTester(model_name, mode)
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
+
     if mode == "eval":
-        logits = results.logits
+        logits = process_batched_logits(results.logits, batch_size)
         predicted_class_id = logits.argmax().item()
         predicted_label = tester.model.config.id2label[predicted_class_id]
 

tests/models/albert/test_albert_token_classification.py

@@ -3,7 +3,7 @@
 from transformers import AutoTokenizer, AlbertForTokenClassification
 import torch
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
 
 
 class ThisTester(ModelTester):
@@ -27,15 +27,27 @@ def _load_inputs(self):
         pytest.param("albert/albert-base-v2", marks=pytest.mark.converted_end_to_end),
     ],
 )
-def test_albert_token_classification(record_property, model_name, mode):
+def test_albert_token_classification(record_property, model_name, mode, get_batch_size):
     record_property("model_name", f"{model_name}-classification")
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
-        logits = results.logits
+        if batch_size is not None:
+            results.logits = results.logits.squeeze(0) # Temporary fix, not the neatest solution
+
+        logits = process_batched_logits(results.logits, batch_size).unsqueeze(0)
+        if batch_size is None:
+            logits = logits.squeeze(0) # Adjust dimensions to account for batch reshaping ^
+
         predicted_token_class_ids = logits.argmax(-1)
 
         # Note that tokens are classified rather then input words which means that

tests/models/autoencoder_conv/test_autoencoder_conv_v2.py

@@ -4,7 +4,8 @@
 import torchvision.transforms as transforms
 from datasets import load_dataset
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
+
 
 
 class ConvAE(torch.nn.Module):
@@ -71,13 +72,19 @@ def _load_inputs(self):
     "mode",
     ["train", "eval"],
 )
-def test_autoencoder_conv_v2(record_property, mode):
+def test_autoencoder_conv_v2(record_property, mode, get_batch_size):
     model_name = f"Autoencoder (conv)"
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
         print("Output: ", results)

tests/models/autoencoder_linear/test_autoencoder_linear.py

@@ -4,7 +4,8 @@
 import torchvision.transforms as transforms
 from datasets import load_dataset
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
+
 
 
 class LinearAE(torch.nn.Module):
@@ -80,13 +81,18 @@ def _load_inputs(self):
     "mode",
     ["train", pytest.param("eval", marks=pytest.mark.converted_end_to_end)],
 )
-def test_autoencoder_linear(record_property, mode):
+def test_autoencoder_linear(record_property, mode, get_batch_size):
     model_name = "Autoencoder (linear)"
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size)
 
     if mode == "eval":
         print("Output: ", results)

tests/models/beit/test_beit_image_classification.py

@@ -3,7 +3,7 @@
 import requests
 import pytest
 import torch
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
 
 
 class ThisTester(ModelTester):
@@ -33,12 +33,18 @@ def get_results_train(self, model, inputs, outputs):
 
 @pytest.mark.parametrize("mode", ["train", "eval"])
 @pytest.mark.parametrize("model_name", ["microsoft/beit-base-patch16-224", "microsoft/beit-large-patch16-224"])
-def test_beit_image_classification(record_property, model_name, mode):
+def test_beit_image_classification(record_property, model_name, mode, get_batch_size):
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
         logits = results.logits

tests/models/bert/test_bert.py

@@ -3,7 +3,7 @@
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
 
 
 class ThisTester(ModelTester):
@@ -35,13 +35,19 @@ def _load_inputs(self):
     ["eval"],
 )
 @pytest.mark.converted_end_to_end
-def test_bert(record_property, mode):
+def test_bert(record_property, mode, get_batch_size):
     model_name = "BERT"
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
         # Helper function to decode output to human-readable text

tests/models/bloom/test_bloom.py

@@ -3,7 +3,8 @@
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
+
 
 
 class ThisTester(ModelTester):
@@ -33,13 +34,19 @@ def _load_inputs(self):
     ["eval"],
 )
 @pytest.mark.converted_end_to_end
-def test_bloom(record_property, mode):
+def test_bloom(record_property, mode, get_batch_size):
     model_name = "Bloom"
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
         # Helper function to decode output to human-readable text

tests/models/distilbert/test_distilbert.py

@@ -1,7 +1,8 @@
 from transformers import DistilBertTokenizer, DistilBertModel
 import torch
 import pytest
-from tests.utils import ModelTester
+from tests.utils import ModelTester, validate_batch_size, process_batched_logits, batch_object_inputs
+
 
 
 class ThisTester(ModelTester):
@@ -22,12 +23,18 @@ def _load_inputs(self):
 )
 @pytest.mark.converted_end_to_end
 @pytest.mark.parametrize("model_name", ["distilbert-base-uncased"])
-def test_distilbert(record_property, model_name, mode):
+def test_distilbert(record_property, model_name, mode, get_batch_size):
     record_property("model_name", model_name)
     record_property("mode", mode)
 
+    batch_size = get_batch_size
+    if batch_size is not None:
+        batch_size = int(batch_size)
+    validate_batch_size(batch_size)
+
     tester = ThisTester(model_name, mode)
-    results = tester.test_model()
+    results = tester.test_model(batch_size=batch_size)
+    batch_object_inputs(tester, batch_size) # This is necessary to avoid shape mismatch errors in tester processing
 
     if mode == "eval":
         print(f"Model: {model_name} | Input: {tester.text} | Output: {results}")