#0: Added a naive batching mechanism for some models in the tests folder #740
Conversation
| if mode == "eval": | ||
| # retrieve index of [MASK] | ||
|
|
||
| results.logits = process_batched_logits(results.logits, batch_size) | ||
| #print(results.logits.shape) |
For testing purposes
| @@ -196,8 +202,8 @@ def compile_and_run(device, reset_torch_dynamo, request): | |||
| model_name, option._aten_fx_graphs, option._out_fx_graphs, option._all_inputs | |||
| ) | |||
|
|
|||
| if len(option._out_fx_graphs) > 0: | |||
| option._out_fx_graphs[0].print_tabular() | |||
| # if len(option._out_fx_graphs) > 0: | |||
There was a problem hiding this comment.
restore?
Will restore
tests/conftest.py
Outdated
| @@ -70,6 +71,11 @@ def device(): | |||
| ttnn.close_device(device) | |||
|
|
|||
|
|
|||
| @pytest.fixture(scope="session") | |||
| def get_batch_size(request): | |||
There was a problem hiding this comment.
Returned value is a value, not a function, so the name starting with "get" is a bit weird here.
Wdyt?
There was a problem hiding this comment.
Maybe just "batch_size"
I agree, looks cleaner
| if mode not in ["train", "eval"]: | ||
| raise ValueError(f"Current mode is not supported: {mode}") | ||
| self.model_name = model_name | ||
| self.mode = mode | ||
| self.model = self._load_model() | ||
| self.inputs = self._load_inputs() | ||
| self.batch_size = batch_size | ||
| self.validate_batch_size() | ||
| self.batch_inputs() |
There was a problem hiding this comment.
Should we instead make load_inputs() be responsible for that?
Maybe _load_inputs() must accept the batch_size parameter or respect the batch_size property of the class.
Just a thought. Open to discuss.
There was a problem hiding this comment.
One of the things is that we are making a shortcut right now, just cloning input. But ideally, we should have N different inputs batched into a single run.
We can have a shortcut now, but from test arch point of view, it is great if its easy to later make an improvement for a given model.
There was a problem hiding this comment.
Should we instead make
load_inputs()be responsible for that?Maybe
_load_inputs()must accept thebatch_sizeparameter or respect the batch_size property of the class. Just a thought. Open to discuss.
I think it might read a bit cleaner, but I think this way is better because there is less code, so its easier to change in the future.
Added a batching flag for the some models in the test folder, repeats the same input tensor across a new dimension. Also misc supporting functions for that.