optimized train time for a use case of small samples and large batch #268
Conversation
| has_error = True | ||
| has_missing_values = True | ||
| if self._raise_error_key_missing: | ||
| raise Exception(f"Error: key {key} does not exist in sample {index}: {sample}") | ||
| else: | ||
| value = None | ||
| else: |
There was a problem hiding this comment.
removing to optimize the running time.
Detecting NaNs and more could move to an optional op (in the data pipeline)
fuse/utils/ndict.py
Outdated
| @@ -228,7 +246,7 @@ def pop(self, key: str) -> Any: | |||
| del self[key] | |||
| return res | |||
|
|
|||
| def indices(self, indices: Union[torch.Tensor, numpy.ndarray]) -> dict: | |||
| def indices(self, indices: Optional[numpy.ndarray]) -> dict: | |||
There was a problem hiding this comment.
self review: will remove the optional here
| @@ -68,9 +68,11 @@ def __call__(self, samples: List[Dict]) -> Dict: | |||
| batch_dict = NDict() | |||
|
|
|||
| # collect all keys | |||
| keys = self._collect_all_keys(samples) | |||
| if self._keep_keys: | |||
There was a problem hiding this comment.
In the description we say: "missing keep_keys are skipped." , I think that now we won't do that.
Could it be an issue? If a user specifies to keep a key that doesn't exist?
There was a problem hiding this comment.
Good catch Now we will throw an error in such a case. I will update the comment.
| batch_size = len(batch["data.sample_id"]) | ||
| else: | ||
| batch_size = None | ||
|
|
||
| if batch_size is None: | ||
| keys = batch.keys() | ||
|
|
||
| for key in keys: | ||
| if isinstance(batch[key], torch.Tensor): |
There was a problem hiding this comment.
Why there are two different loops? one for each case - torch.Tensor, (np.ndarray, list)
If I'm not missing something we can check for the two cases in the same loop.
There was a problem hiding this comment.
Because, first I want to look for tensor (trust it more), and if I can't find one then my second choice is (np.ndarray, list)
|
|
||
| return all_keys | ||
| @staticmethod | ||
| def _flatten_static(item: Union[dict, Any], prefix: str, flat_dict: dict) -> None: |
|
|
||
| def keypaths(self) -> List[str]: | ||
| """ | ||
| returns a list of keypaths (i.e. "a.b.c.d") to all values in the nested dict | ||
| """ | ||
| return list(self.flatten().keys()) | ||
| return NDict._keypaths_static(self._stored, None) |
There was a problem hiding this comment.
Why not using the same paradigm as before? Just calling flatten()?
The two static functions has a lot in common
There was a problem hiding this comment.
Mostly cause I don't want the overhead of creating a dictionary and extracting the keys,
Optimize nicely the train running time in such a use case (small samples and large batch size)
Leaving ndict optimization for the future