facebookresearch / facebookresearch/detectron2

A simple trick for a fully deterministic ROIAlign, and thus MaskRCNN training and inference

Open
#4,723 6 comments 9 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
34.7k
Forks
7.9k
PR merge metrics
No merged PRs in 30d

Description

### Non-determinism of MaskRCNN

There have been a lot of discussions and inquiries in this repo about a fully deterministic MaskRCNN e.g. #4260, #3203 , #2615, #2480, and also on other detection repositories (e.g. MMDetection [here](https://github.com/open-mmlab/mmdetection/issues/2558) and [here](https://github.com/open-mmlab/mmdetection/issues/2773) and also torchvision [here](https://github.com/pytorch/vision/issues/1058)). Unfortunately, even after seeding everything and setting Pytorch's [deterministic flags](https://pytorch.org/docs/stable/notes/randomness.html), results are still non-repeatable.

It boils down to the fact that some of the used Pytorch / torchvision ops doesn't have a deterministic GPU implementation (most notably, due to using `atomicAdd` in the backward pass). So, the only solution is to train for as long as possible to reduce variance in the results. It is worth noting that not only training, but also evaluation (see #2480) of MaskRCNN (and actually most detectron2 models) is not deterministic

Based on the minimal example in #4260, I made an analysis on the ops used for MaskRCNN and found that the main reason of non-determinism is the backward pass of `ROIAlign` (see [here](https://github.com/pytorch/vision/blob/main/torchvision/csrc/ops/cuda/roi_align_kernel.cu#L303)).

### Proposed solution

I am here proposing a simple trick that makes `ROIAlign` practically fully reproducible, without touching the cuda kernel!! it introduces trivial additional memory and computation. It can be summarized as:
- Truncate the input to a smaller datatype, this gives a starting point with a very small number of significand bits used
- Then, cast to a larger data-type just before doing the computations that involve `atomicAdd`

In terms of code, this is translated to simply modifying this [function call](https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/roi_align.py#L58) to

```python
return roi_align(
input.half().double(),
rois.half().double(),
self.output_size,
self.spatial_scale,
self.sampling_ratio,
self.aligned,
).to(dtype=input.dtype)
```

### Test

The conversion to `double` results in a trivial increase in memory & computation, but performing it after the truncation, significantly increases reproducibility.

This solution was tested and found fully deterministic (losses values, and evaluation results on COCO) upto tens of thousands of steps (using same code as in #4260) for:
- MaskRCNN based on ResNet-50 bakbone
- MaskRCNN based on ResNeXt-101 bakbone
- Wide range of batch sizes
- Mixed-precision training
- Single and Multi-GPU training
- A100's & V100's

### Note on A100

Ampere by default uses TF32 format for tensor-core computations, which means that the above truncation is done implicitly! so on Ampere based devices it is enough just to cast to double, i.e.

```python
return roi_align(
input.double(),
rois.double(),
self.output_size,
self.spatial_scale,
self.sampling_ratio,
self.aligned,
).to(dtype=input.dtype)
```

Note: This is the default mode for PyTorch, but if TF32 is disabled for some reason (i.e. `torch.backends.cudnn.allow_tf32 = False`) then the above truncation with `.half()` is still necessary

### Note
- This solution was tested and found to work well for other non-deterministic Pytorch ops, including: [`F.interpolate`](https://pytorch.org/docs/stable/generated/torch.nn.functional.interpolate.html) and [`F.grid_sample`](https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html)
- This is not a general solution to the problem of random-order reproducible floating point summation, but a practical mitigation that works well for this setup / scenario
- At least in theory, this should work even better if applied inside the kernel right before `atomicAdd`
- The only alternative currently is training each experiment for very long, which isn't practical in many setups, and still isn't fully reproducible

Would love to hear what people think about this!
@ppwwyyxx @fmassa

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.