Add missing DEFAULT_TRAIN_BSIZE field

pytorch · Jan 13, 2025 · e6fde3b · e6fde3b
1 parent 9028c93
commit e6fde3b
Show file tree

Hide file tree

Showing 11 changed files with 11 additions and 0 deletions.
diff --git a/torchbenchmark/models/doctr_det_predictor/__init__.py b/torchbenchmark/models/doctr_det_predictor/__init__.py
@@ -13,6 +13,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.DETECTION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

diff --git a/torchbenchmark/models/doctr_reco_predictor/__init__.py b/torchbenchmark/models/doctr_reco_predictor/__init__.py
@@ -13,6 +13,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.DETECTION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

diff --git a/torchbenchmark/models/llama/__init__.py b/torchbenchmark/models/llama/__init__.py
@@ -12,6 +12,7 @@
 class Model(BenchmarkModel):
     task = NLP.LANGUAGE_MODELING
     DEFAULT_EVAL_BSIZE = 32
+    DEFAULT_TRAIN_BSIZE = 32
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):
         super().__init__(

diff --git a/torchbenchmark/models/pyhpc_equation_of_state/__init__.py b/torchbenchmark/models/pyhpc_equation_of_state/__init__.py
@@ -40,6 +40,7 @@ class Model(BenchmarkModel):
     # Source: https://github.com/dionhaefner/pyhpc-benchmarks/blob/650ecc650e394df829944ffcf09e9d646ec69691/run.py#L25
     # Pick data point: i = 20, size = 1048576
     DEFAULT_EVAL_BSIZE = 1048576
+    DEFAULT_TRAIN_BSIZE = 1048576
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

diff --git a/torchbenchmark/models/pyhpc_isoneutral_mixing/__init__.py b/torchbenchmark/models/pyhpc_isoneutral_mixing/__init__.py
@@ -131,6 +131,7 @@ class Model(BenchmarkModel):
     # Source: https://github.com/dionhaefner/pyhpc-benchmarks/blob/650ecc650e394df829944ffcf09e9d646ec69691/run.py#L25
     # Pick data-point when i = 20, size = 1048576
     DEFAULT_EVAL_BSIZE = 1048576
+    DEFAULT_TRAIN_BSIZE = 1048576
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

diff --git a/torchbenchmark/models/pyhpc_turbulent_kinetic_energy/__init__.py b/torchbenchmark/models/pyhpc_turbulent_kinetic_energy/__init__.py
@@ -129,6 +129,7 @@ class Model(BenchmarkModel):
     # Source: https://github.com/dionhaefner/pyhpc-benchmarks/blob/650ecc650e394df829944ffcf09e9d646ec69691/run.py#L25
     # Pick data-point when i = 20, size = 1048576
     DEFAULT_EVAL_BSIZE = 1048576
+    DEFAULT_TRAIN_BSIZE = 1048576
     ALLOW_CUSTOMIZE_BSIZE = False
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 

diff --git a/torchbenchmark/models/sam/__init__.py b/torchbenchmark/models/sam/__init__.py
@@ -17,6 +17,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.SEGMENTATION
     DEFAULT_EVAL_BSIZE = 32
+    DEFAULT_TRAIN_BSIZE = 32
 
     def __init__(self, test, device, batch_size=1, extra_args=[]):
         super().__init__(

diff --git a/torchbenchmark/models/sam_fast/__init__.py b/torchbenchmark/models/sam_fast/__init__.py
@@ -18,6 +18,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.SEGMENTATION
     DEFAULT_EVAL_BSIZE = 32
+    DEFAULT_TRAIN_BSIZE = 32
 
     def __init__(self, test, device, batch_size=1, extra_args=[]):
         super().__init__(

diff --git a/torchbenchmark/models/simple_gpt/__init__.py b/torchbenchmark/models/simple_gpt/__init__.py
@@ -11,6 +11,7 @@
 class Model(BenchmarkModel):
     task = NLP.GENERATION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
 
     def validate_environment(self):
         if not torch.cuda.is_available() or "cuda" not in self.device:

diff --git a/torchbenchmark/models/simple_gpt_tp_manual/__init__.py b/torchbenchmark/models/simple_gpt_tp_manual/__init__.py
@@ -9,6 +9,7 @@
 class Model(BenchmarkModel):
     task = NLP.GENERATION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
 
     def validate_environment(self):
         if not torch.cuda.is_available() or "cuda" not in self.device:

diff --git a/torchbenchmark/util/framework/huggingface/model_factory.py b/torchbenchmark/util/framework/huggingface/model_factory.py
@@ -140,6 +140,7 @@ def __init__(self):
 class HuggingFaceGenerationModel(HuggingFaceModel):
     task = NLP.GENERATION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
 
     """
     Instead of just running __call__ on the model, use generate to generate