• Vasiliy Kuznetsov's avatar
    vision classification QAT tutorial: fix for DDP (redo) (#2230) · 7ed3950e
    Vasiliy Kuznetsov authored
    Summary:
    
    Redo of https://github.com/pytorch/vision/pull/2191
    
    Makes the classification QAT tutorial not crash when used
    with DDP. There were two issues:
    
    1. the model was moved to GPU before the observers were added, and they
    are created on CPU. In the context of this repo, the fix is to finalize
    the model before moving to GPU. We can potentially follow up with a
    better error message in the future, in a separate PR.
    2. the QAT conversion was running on the DDP'ed model, which had various
    problems. The fix is to unwrap the model from DDP before cloning it for
    evaluation.
    
    There is still work to do on verifying that BN is working correctly in
    QAT + DDP, but saving that for a separate PR.
    
    Test Plan:
    
    ```
    python -m torch.distributed.launch --use_env references/classification/train_quantization.py --data-path {path_to_imagenet_1k} --output_dir {output_dir}
    ```
    
    Reviewers:
    
    Subscribers:
    
    Tasks:
    
    Tags:
    7ed3950e
train_quantization.py 10.4 KB