Applying Apex amp to DETR

kkjh0723: <p>I’m trying to applying <code>apex.amp</code> to  recent detection transformer (DETR) code (<a href="https://github.com/facebookresearch/detr" rel="nofollow noopener">link</a>)<br>
What I’m not sure is where to put <code>amp.initialize</code><br>
Here are lines from the <code>main.py</code> of DETR where model and optimizer are declared (from <a href="https://github.com/facebookresearch/detr/blob/10a2c759454930813aeac7af5e779f835dcb75f5/main.py#L121" rel="nofollow noopener">line#121</a>)</p>
<pre><code class="lang-auto">    model, criterion, postprocessors = build_model(args)
    model.to(device)

    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        model_without_ddp = model.module
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('number of params:', n_parameters)

    param_dicts = [
        {"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
        {
            "params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
            "lr": args.lr_backbone,
        },
    ]
    optimizer = torch.optim.AdamW(param_dicts, lr=args.lr,
                                  weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
</code></pre>
<p>According to the apex imagenet example, <code>amp.initialize</code> should place before wrapping model with DDP. I can move DDP wrapping part to after the <code>lr_scheduler</code> like below,</p>
<pre><code class="lang-auto">    model, criterion, postprocessors = build_model(args)
    model.to(device)

    model_without_ddp = model  

    param_dicts = [
        {"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
        {
            "params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
            "lr": args.lr_backbone,
        },
    ]
    optimizer = torch.optim.AdamW(param_dicts, lr=args.lr,
                                  weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)

    # Apex amp
    model, optimizer = amp.initialize(model, optimizer,
                                      opt_level=args.opt_level,
                                      keep_batchnorm_fp32=args.keep_batchnorm_fp32,
                                      loss_scale=args.loss_scale
                                      )
    
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        model_without_ddp = model.module
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('number of params:', n_parameters)
</code></pre>
<p>But I wonder if I can pass <code>model_without_ddp</code> to <code>amp.initialize</code> after the end of above code without moving the DDP wrapping lines.</p>
<pre><code class="lang-auto">...
    optimizer = torch.optim.AdamW(param_dicts, lr=args.lr,
                                  weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
    
    # Apex amp
    model, optimizer = amp.initialize(model_without_ddp, optimizer,
                                      opt_level=args.opt_level,
                                      keep_batchnorm_fp32=args.keep_batchnorm_fp32,
                                      loss_scale=args.loss_scale
                                      )

</code></pre>

ptrblck: <p>We recommended to setup the mixed precision training before wrapping the model into DDP.<br>
Wouldn’t this work for you or what the reason you would like to setup DDP before amp?</p>
<p>Also, note that we recommend trying out the native amp implementation using the nightly binaries, as explained <a href="https://discuss.pytorch.org/t/torch-cuda-amp-vs-nvidia-apex/74994/3">here</a>.</p>

kkjh0723: <p>Thanks for your help.<br>
Not a important reason. I thought it might be possible to add amp without any changes of original code if the second option is valid.</p>
<p>I also considered to use native amp in nightly version.<br>
But I was not sure it would be better to use the nightly version or wait until stable version is released.<br>
Maybe I should try native amp also.</p>

kkjh0723: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, while I’m trying 1st option using apex, I got following error.</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "main_apex.py", line 269, in &lt;module&gt;
    main(args)
  File "main_apex.py", line 219, in main
    args.clip_max_norm)
  File "/home/detr/engine_apex.py", line 37, in train_one_epoch
    loss_dict = criterion(outputs, targets)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/detr/models/detr.py", line 225, in forward
    indices = self.matcher(outputs_without_aux, targets)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/home/detr/models/matcher.py", line 71, in forward
    cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)
  File "/usr/local/lib/python3.7/dist-packages/torch/functional.py", line 732, in cdist
    return _VF.cdist(x1, x2, p, None)
RuntimeError: "cdist_cuda" not implemented for 'Half'
</code></pre>
<p>It seems <code>torch.cdist</code> function is not working with half precision.<br>
Is there any workaround?<br>
Can this issue be resolved in nightly version?</p>

ptrblck: <p>You could decorate the method with <a href="https://github.com/NVIDIA/apex/blob/master/apex/amp/README.md#function-annotation"><code>@amp.float_function</code></a> (maybe you would have to wrap it into another helper method), which should disable the casting.</p>
<aside class="quote no-group" data-username="kkjh0723" data-post="4" data-topic="86984">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/kkjh0723/48/3571_2.png" class="avatar"> kkjh0723:</div>
<blockquote>
<p>Can this issue be resolved in nightly version?</p>
</blockquote>
</aside>
<p>This might be the case. If you see that it’s not working in native amp, let us know, please.</p>

kkjh0723: <p>I’m working on native amp and found that above issue is not occurred because <code>torch.cdist</code> is included in Ops that autocast to <code>float32</code>.</p>
<p>However, I faces another issue. The <code>targets</code> in the code which is a <code>list</code> of <code>dict</code> of <code>tensors</code> (<a href="https://github.com/facebookresearch/detr/blob/10a2c759454930813aeac7af5e779f835dcb75f5/engine.py#L30" rel="nofollow noopener">line</a>) and it passes to the <code>criterion</code>.<br>
It seems <code>autocast</code> cannot handle this type of input.<br>
What is the best way to make it work with the <code>autocast</code>?<br>
Currently, I disable <code>autocast</code> at each loss and cast the outputs and targets to <code>float</code> following <a href="https://pytorch.org/docs/master/notes/amp_examples.html#autocast-and-custom-autograd-functions" rel="nofollow noopener">this</a>.</p>
<p>It seems loss functions are mostly calculated with <code>float32</code>. Is it right?<br>
I wonder if there is a <code>float_function</code> like decorator for the method (that you informed above) in native amp so that I can decorate the loss function such as <a href="https://github.com/facebookresearch/detr/blob/10a2c759454930813aeac7af5e779f835dcb75f5/util/box_ops.py#L40" rel="nofollow noopener">this</a>.</p>

Jacobew: <p>Hi, did you figure it out? I’ve encountered the same issue when applying native apex to DETR.<br>
I’ve tried decorating SetCriterion.foward() with amp.float_function but the errors remain.</p>