Apex: memory gains in detectron2

username: soldierofhell
cooked: <p>Hi, I modified <a href="https://github.com/facebookresearch/detectron2" rel="nofollow noopener">detectron2</a> <a href="https://github.com/facebookresearch/detectron2/blob/master/configs/Misc/cascade_mask_rcnn_X_152_32x8d_FPN_IN5k_gn_dconv.yaml" rel="nofollow noopener">Mask X152</a> model for apex training. Besides standard “three lines of code” change in trainer’s class, I’ve decorated roialign and deconv custom layers to force 32-bit for them.<br>
My main goal was to reduce memory, but unfortunately it didn’t happen.  Although after reading brief <a href="https://nvidia.github.io/apex/amp.html#o1-mixed-precision-recommended-for-typical-use" rel="nofollow noopener">description of recommended opt_level=O1</a> it’s not clear to me, if I should expect them at all. Seems like only computations are done in 16bit (after casting), but “storage” is mostly 32bit.</p>
<p>Should I try to adjust code to opt_level=O2?</p>

username: ptrblck
cooked: <p>You could still save memory with <code>opt_level='O1'</code>, as activations might be stored in FP16 as seen in this dummy code snippet:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

import torchvision.models as models

from apex import amp

use_amp = True

model = models.resnet50()
model.cuda()


model.avgpool.register_forward_hook(lambda m, x, y: print(x[0].type(), y[0].type()))
x = torch.randn(1, 3, 224, 224).cuda()

torch.cuda.synchronize()
print(torch.cuda.max_memory_allocated()/1024**2)

if use_amp:
    model = amp.initialize(model, opt_level='O1')
output = model(x)

torch.cuda.synchronize()
print(torch.cuda.max_memory_allocated()/1024**2)
</code></pre>
<p>besides the speedup using TensorCores.</p>

username: soldierofhell
cooked: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a>, I thought about using forward_hook. Indeed starting from ‘O1’ we see torch.cuda.HalfTensor, but…<br>
Here are the memory results for O0, O1 and O2 for RTX 2070 SUPER. How can we explain higher memory usage by O1? (same one can observe on google colab)</p>
<pre><code class="lang-auto">O0
check if max_memory strarts from zero:  0.0
memory after model and data are loaded:  98.30224609375
Selected optimization level O0:  Pure FP32 training.

Defaults for this optimization level are:
enabled                : True
opt_level              : O0
cast_model_type        : torch.float32
patch_torch_functions  : False
keep_batchnorm_fp32    : None
master_weights         : False
loss_scale             : 1.0
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O0
cast_model_type        : torch.float32
patch_torch_functions  : False
keep_batchnorm_fp32    : None
master_weights         : False
loss_scale             : 1.0
memory after forward:  181.4619140625
</code></pre>
<pre><code class="lang-auto">O1
check if max_memory strarts from zero:  0.0
memory after model and data are loaded:  98.30224609375
Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
memory after forward:  192.01171875
</code></pre>
<pre><code class="lang-auto">O2
check if max_memory strarts from zero:  0.0
memory after model and data are loaded:  98.30224609375
Selected optimization level O2:  FP16 training with FP32 batchnorm and FP32 master weights.

Defaults for this optimization level are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
memory after forward:  98.892578125
</code></pre>
<p>Modified code snippet:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

import torchvision.models as models

from apex import amp

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

opt_level = 'O1'

print(opt_level)

seed = 123
torch.manual_seed(seed)

torch.cuda.reset_max_memory_allocated()

print('check if max_memory strarts from zero: ', torch.cuda.max_memory_allocated()/1024**2)

model = models.resnet50()
model.cuda()

#model.avgpool.register_forward_hook(lambda m, x, y: print(x[0].type(), y[0].type()))
x = torch.randn(1, 3, 224, 224).cuda()

torch.cuda.synchronize()
print('memory after model and data are loaded: ', torch.cuda.max_memory_allocated()/1024**2)

model = amp.initialize(model, opt_level=opt_level)
output = model(x)

torch.cuda.synchronize()
print('memory after forward: ', torch.cuda.max_memory_allocated()/1024**2)
</code></pre>

username: ptrblck
cooked: <p>Could you also check the allocated memory via <code>torch.cuda.memory_allocated()</code>?</p>

username: soldierofhell
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, exactly the same, so max is equal current. This is from pytorch ‘1.4.0+cu100’</p>

username: ptrblck
cooked: <p>Thanks for the information.<br>
Since you’ve set <code>torch.backends.cudnn.deterministic = True</code>, the applied algorithms might have a different memory requirement for different data types, which seems to be higher in <code>O1</code> than in <code>O0</code> for this use case.</p>

username: Walid_Ahmed_pyTorch
cooked: <p><a class="mention" href="/u/soldierofhell">@soldierofhell</a><br>
I am trying to use O2   but  roialign  is giving me the following error<br>
"File “/opt/conda/lib/python3.7/site-packages/torchvision/ops/roi_align.py”, line 45, in roi_align<br>
sampling_ratio, aligned)<br>
RuntimeError: Expected tensor for argument <span class="hashtag">#1</span> ‘input’ to have the same type as tensor for argument <span class="hashtag">#2</span> ‘rois’; but type torch.cuda.HalfTensor does not equal torch.cuda.FloatTensor (while checking arguments for ROIAlign_forward_cuda)<br>
"<br>
Can you please guide me how to resolve this error?</p>
<p>Thanks<br>
Walid</p>

username: ptrblck
cooked: <p>I would recommend to use the native <a href="https://pytorch.org/docs/stable/amp.html">mixed precision package</a>, which should support the current <code>torchvision</code> models with their custom extensions.</p>

username: Walid_Ahmed_pyTorch
cooked: <p>Thanks a lot.<br>
This native api gave me the following error:<br>
detectron2/modeling/poolers.py", line 249, in forward<br>
output[inds] = pooler(x[level], pooler_fmt_boxes_level)<br>
RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and Half for the source.</p>
<p>Can you please help me how to resolve it.</p>
<p>Also how can I use O1 and O2 levels too  like we   do in <a href="https://nvidia.github.io/apex/amp.html" rel="noopener nofollow ugc">https://nvidia.github.io/apex/amp.html</a>?</p>
<p>Walid</p>

username: ptrblck
cooked: <p>Native amp doesn’t support different <code>opt_level</code>s and is similar to <code>O1</code>.<br>
Could you install the latest nightly binaries and check, if the error is still raised, please?</p>

username: Walid_Ahmed_pyTorch
cooked: <p>Thanks, I will give it a try but my main aim is o2. I already got o1 working with nvidia Apex</p>