Introduce Hyrax's own auto_model function to remove DataParallel#996
Introduce Hyrax's own auto_model function to remove DataParallel#996SamSandwich07 wants to merge 11 commits into
Conversation
| torch.nn.Module | ||
| """ | ||
| wrapped_model = idist.auto_model(model, sync_bn=sync_bn, **kwargs) | ||
| if type(wrapped_model) is DataParallel: |
There was a problem hiding this comment.
We could simply make this a RuntimeError. Thoughts?
There was a problem hiding this comment.
No, that wouldn't work. idist.auto_model is going to make it DataParallel whenever we run Hyrax in an environment with multiple GPUs (like Gondor) but doesn't use idist.Parallel (i.e. set distributed to false or run a verb other than train). If we raise an error here then Hyrax will essentially always raise an error on Gondor and the like.
|
|
||
| Args: | ||
| model: model to adapt. | ||
| sync_bn: if True, applies `torch convert_sync_batchnorm`_ to the model for native torch |
There was a problem hiding this comment.
Do we know we will need to change sync_bn from callers of this function?
There was a problem hiding this comment.
I guess I'm unclear why we're promoting it out of **kwargs in this interface.
If we need it set to false consistently, then we can put that, refactor when usage forces it?
There was a problem hiding this comment.
We are the only callers of idist.auto_model, and in our code we never call it with any arguments other than the torch model, so the answer for now is "we never change the argument".
There was a problem hiding this comment.
As for why we promote it out of **kwargs, it's because idist.auto_model keeps it outside of **kwargs too, so I just used the same input signature (below taken from ignite docs):
def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Module:
Introduce a helper function
_auto_modeltopytorch_ignite.pyto replaceidist.auto_modeland removeDataParallel. In other words, models will never be wrapped inDataParallelanymore; they will either be unwrapped or wrapped inDistributedDataParallel.