Continuing cleanup. #76
Conversation
znado commented May 8, 2022
znado commented - refactoring dataset iterators to return Dict[str, Any] instead of a tuple/triplet of inputs/labels/masks
- cleaning up some of the WMT code
- standardizing param shape/type utilities
- small cleanups/fixes to get all reference submissions to work with random inputs for 1 train and 2 eval steps (takes ~272s total to test)
| MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
There was a problem hiding this comment.
Looks really awesome!
I added a few comments and nits. I think, after merging #66 and addressing these comments, this is ready to merge.
| # flax.linen.Embed names the embedding parameter "embedding" | ||
| # https://github.com/google/flax/blob/main/flax/linen/linear.py#L604. | ||
| elif name == 'embedding': | ||
| param_types_dict[name] = spec.ParameterType.EMBEDDING |
There was a problem hiding this comment.
Is it problematic that JAX gets more information than PyTorch?
There was a problem hiding this comment.
yes. is the way I added the right way to check the name of an nn.Embedding layer in Pytorch? we should add a test to make sure (for jax and pytorch) these are properly set (at least, for some selection of params). I added a TODO to my list of things to follow up with.
| | ||
| if is_train: | ||
| dataloader = cycle(dataloader) | ||
| dataloader = cycle(dataloader) |
There was a problem hiding this comment.
Just curious why you switched to also cycle the evaluation datasets?
There was a problem hiding this comment.
I added it so that we can iterate through them across multiple evals. I'm not sure if this is the best way to do it.
algorithmic_efficiency/workloads/librispeech/librispeech_pytorch/workload.py Outdated Show resolved Hide resolved
2 participants
