Apex with `torch.optim.swa_utils`

username: Sayak_Paul
cooked: <p>I am currently trying to using <code>apex</code> with SWA like so:</p>
<pre><code class="lang-python">model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
model = DDP(model)

swa_model = torch.optim.swa_utils.AveragedModel(model)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
swa_start = 6
swa_scheduler = SWALR(optimizer, swa_lr=0.05 * world_size)
</code></pre>
<p>At this line (<code>swa_model = torch.optim.swa_utils.AveragedModel(model)</code>) I am getting the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/home/jupyter/Flood_Comp/starter.py", line 249, in train
    swa_model = torch.optim.swa_utils.AveragedModel(model)
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/optim/swa_utils.py", line 89, in __init__
    self.module = deepcopy(model)
  File "/opt/conda/lib/python3.7/copy.py", line 169, in deepcopy
    rv = reductor(4)
  File "/opt/conda/lib/python3.7/site-packages/apex/parallel/distributed.py", line 271, in __getstate__
    del attrs['self.bucket_streams']
KeyError: 'self.bucket_streams'
</code></pre>
<p>Any pointers on mitigating this would be helpful.</p>

username: ptrblck
cooked: <p><code>apex.amp</code> is deprecated in favor of the native <code>torch.cuda.amp</code> implementation and we recommend to switch to the latter.<br>
More details are given in <a href="https://discuss.pytorch.org/t/torch-cuda-amp-vs-nvidia-apex/74994/3">this post</a>.</p>

username: Sayak_Paul
cooked: <p>Okay. Strangely enough, I am unable to get the <code>apex.amp</code> benefits in <code>torch.cuda.amp</code>. But I look into the post you suggested and see if there’s anything I am missing out on.</p>

username: Sayak_Paul
cooked: <p>After following <a class="mention" href="/u/ptrblck">@ptrblck</a>’s suggestions here’s how my <code>train()</code> function is looking like (consider this to be the launcher expected by <code>torch.multiprocessing.spawn()</code>.</p>
<pre><code class="lang-python">def train(rank, num_epochs, world_size):
    init_process(rank, world_size)
    torch.manual_seed(0)
    
    model = create_model()
    torch.cuda.set_device(rank)
    model.cuda(rank)
    model = DistributedDataParallel(model, device_ids=[rank])
    swa_model = torch.optim.swa_utils.AveragedModel(model)

    learning_rate = 1e-3
    optimizer = torch.optim.Adam(model.parameters(), 
                             lr=learning_rate * world_size)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
    swa_start = 10
    swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, swa_lr=0.05)
    criteria = nn.CrossEntropyLoss()
    scaler = torch.cuda.amp.GradScaler(enabled=True)

    
    train_loader, val_loader = get_dataloader(rank, world_size)
    
    for epoch in range(num_epochs):
    	model.train()
        for batch in train_loader:
            with torch.cuda.amp.autocast(enabled=True):
                image = batch['image'].cuda(rank, non_blocking=True)
                mask = batch['mask'].cuda(rank, non_blocking=True)
                
                pred = model(image)
                
                loss = criteria(pred, mask.unsqueeze(1))

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)
            
        if epoch &gt; swa_start:
            swa_model.update_parameters(model)
            swa_scheduler.step()
        else:
            scheduler.step()
</code></pre>
<p>Things are now working as expected. <a class="mention" href="/u/ptrblck">@ptrblck</a> if I could do anything better to further optimize the performance please let me know.</p>
<p>Majority of the SWA code comes from the <a href="https://pytorch.org/docs/stable/optim.html#stochastic-weight-averaging" rel="noopener nofollow ugc">official docs</a>.</p>

username: Sayak_Paul
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I am currently running into another problem that is closely related. After training with SWA, we need to update the batch norm statistics (<a href="https://pytorch.org/docs/stable/optim.html#stochastic-weight-averaging" rel="noopener nofollow ugc">reference</a>). Since the structure of my dataset is different from what <code>torch.optim.swa_utils.update_bn()</code> expects, I am doing the following inside <code>train()</code> (recall that <code>train()</code> is the launcher I provide to <code>mp.spawn()</code>):</p>
<pre><code class="lang-python">if rank == 0:
    for batch in train_loader: 
        image = batch['image'].cuda(rank, non_blocking=True)
        prediction = swa_model(image)
</code></pre>
<p>This leads to the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "starter.py", line 343, in &lt;module&gt;
    nprocs=WORLD_SIZE, join=True
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 199, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 157, in start_processes
    while not context.join():
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 118, in join
    raise Exception(msg)
Exception: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/home/jupyter/Flood_Comp/starter.py", line 334, in train
    prediction = swa_model(image)
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/optim/swa_utils.py", line 101, in forward
    return self.module(*args, **kwargs)
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 610, in forward
    self._sync_params()
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 1048, in _sync_params
    authoritative_rank,
  File "/home/jupyter/.local/lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 979, in _distributed_broadcast_coalesced
    self.process_group, tensors, buffer_size, authoritative_rank
RuntimeError: [/pytorch/third_party/gloo/gloo/transport/tcp/pair.cc:575] Connection closed by peer [10.138.0.33]:26791
</code></pre>
<p>Anything I am missing out on?</p>

username: ptrblck
cooked: <p>Are you seeing this error using <code>mp.spawn</code> and <code>swa</code> in isolation or only in combination with <code>torch.cuda.amp</code>?</p>

username: Sayak_Paul
cooked: <p>In combination with the three. I did not try out the isolated part.</p>

username: ptrblck
cooked: <p>Could you try to isolate it further, which would help to debug it more? I.e. in particular it would be interesting to see, if your custom <code>mp</code> approach would work with <code>amp</code> or <code>swa</code> in isolation, as this is often causing trouble, if you are not careful.</p>

username: Sayak_Paul
cooked: <p>With just <code>amp</code>, it works fine. This I have verified before.</p>

username: Sayak_Paul
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> any further suggestions?</p>

username: ptrblck
cooked: <p>I’m not familiar with the internals of <code>swa</code> so you could check the <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html">multiprocessing best-practices</a> as well as the docs about <a href="https://pytorch.org/docs/stable/multiprocessing.html#multiprocessing-cuda-sharing-details">sharing CUDA tensors</a>.</p>