Skip to content

Commit

Permalink
IABN2float
Browse files Browse the repository at this point in the history
  • Loading branch information
talrid committed Apr 4, 2020
1 parent 1572729 commit 94f0cc4
Showing 3 changed files with 56 additions and 0 deletions.
44 changes: 44 additions & 0 deletions INPLACE_ABN_TIPS.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
## Some Tips For Working With Inplace-ABN

[Inplace-ABN](https://github.com/mapillary/inplace_abn) have exactly the same fields as
[regular BatchNorm](https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/batchnorm.py):
* module.weight
* module.bias
* module.running_mean
* module.running_var
* module.num_batches_tracked

Therefore, any function that operates on BatchNorm can run on
Inplace-ABN.

However, problems can arise when a logic condition seeks explicitly for
BatchNorm layers only:
```
if isinstance(module, nn.BatchNorm2d):
do_something(module)
```

Anywhere you see a code segement like this, it needs to be replaced with
a condition that include Inplace-ABN:
```
if isinstance(module, nn.BatchNorm2d) or isinstance(module, inplace_abn.InPlaceABN):
do_something(module)
```

##### NVIDIA Apex mixed precision
[NVIDIA Apex](https://github.com/NVIDIA/apex) O0,O1 and O3
mixed-precision options work seemlesly on Inplace-ABN.

For O2 mixed precision, we need to convert manually Inplace-ABN to fp32,
since NVIDIA Apex inner code has explicit 'if' condition for BatchNorm
only.

Conversion can be done easily with the helper function
'IABN2float':
```
if args.use_apex:
model, optimizer = apex.amp.initialize(model, optimizer, opt_level=args.opt_level)
if args.opt_level == 'O2': # IABN needs adjustment for O2
from src.models.tresnet import IABN2float
model = IABN2float(model)
```
5 changes: 5 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -336,6 +336,11 @@ python -m infer.py \
--input_size=224
```

## Tips For Working With Inplace-ABN
See
[INPLACE_ABN_TIPS](https://github.com/mrT23/TResNet/src/INPLACE_ABN_TIPS.md).


## Citation

```
7 changes: 7 additions & 0 deletions src/models/tresnet/tresnet.py
Original file line number Diff line number Diff line change
@@ -7,6 +7,13 @@
from src.models.tresnet.layers.space_to_depth import SpaceToDepthModule
from inplace_abn import InPlaceABN

def IABN2float(module: nn.Module) -> nn.Module:
"If `module` is IABN don't use half precision."
if isinstance(module, InPlaceABN):
module.float()
for child in module.children(): IABN2float(child)
return module

def conv2d_ABN(ni, nf, stride, activation="leaky_relu", kernel_size=3, activation_param=1e-2, groups=1):
return nn.Sequential(
nn.Conv2d(ni, nf, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, groups=groups,

0 comments on commit 94f0cc4

Please sign in to comment.