[optimization] Fix bfloat16 issue and speed up master

docs/source/using_doctr/using_model_export.rst

@@ -11,27 +11,37 @@ Model optimization
 
 This section is meant to help you perform inference with optimized versions of your model.
 
+.. _half-precision:
+
 
 Half-precision
 ^^^^^^^^^^^^^^
 
 **NOTE:** We support half-precision inference for PyTorch models only on **GPU devices**.
 
-Half-precision (or FP16) is a binary floating-point format that occupies 16 bits in computer memory.
+Half-precision formats occupy 16 bits in computer memory instead of the 32 bits used by
+single-precision (FP32). Two formats are supported:
+
+- **BF16** (``bfloat16``): keeps the same exponent range as FP32 with a reduced mantissa.
+- **FP16** (``float16``): higher precision than BF16, but a much narrower dynamic range.
 
 Advantages:
 
 - Faster inference
 - Less memory usage
 
+We recommend **BF16 over FP16**. Because it retains the full FP32 exponent range, BF16 is far
+less prone to overflow and underflow. BF16 requires an Ampere-generation GPU or newer
+(compute capability 8.0+); on older hardware, use FP16 instead.
+
 .. code:: python3
 
     import torch
     predictor = ocr_predictor(
         reco_arch="crnn_mobilenet_v3_small",
         det_arch="linknet_resnet34",
         pretrained=True
-    ).cuda().half()
+    ).cuda().bfloat16()  # or .half() for FP16
     res = predictor(doc)
 
 

docs/source/using_doctr/using_models.rst

@@ -506,7 +506,8 @@ The same approach applies to all standalone predictors:
 * `layout_predictor`
 
 Just create the predictor instance and move it to the appropriate device.
-To enable **half-precision inference**, you can append `.half()` after moving the predictor to the device.
+To enable **half-precision inference**, append `.bfloat16()` after moving the predictor to the
+device -- or `.half()` for FP16, though BF16 is preferred (see :ref:`half-precision` for details)
 
 
 What should I do with the output?

doctr/models/classification/predictor/pytorch.py

@@ -56,7 +56,8 @@ def forward(
         predicted_batches = [self.model(batch) for batch in processed_batches]
         # confidence
         probs = [
-            torch.max(torch.softmax(batch, dim=1), dim=1).values.cpu().detach().numpy() for batch in predicted_batches
+            torch.max(torch.softmax(batch.float(), dim=1), dim=1).values.cpu().detach().numpy()
+            for batch in predicted_batches
         ]
         # Postprocess predictions
         predicted_batches = [out_batch.argmax(dim=1).cpu().detach().numpy() for out_batch in predicted_batches]

doctr/models/detection/differentiable_binarization/base.py

@@ -243,7 +243,7 @@ def draw_thresh_map(
         ys: np.ndarray = np.broadcast_to(np.linspace(0, height - 1, num=height).reshape(height, 1), (height, width))
 
         # Compute distance map to fill the padded polygon
-        distance_map = np.zeros((polygon.shape[0], height, width), dtype=polygon.dtype)
+        distance_map = np.zeros((polygon.shape[0], height, width), dtype=np.float32)
         for i in range(polygon.shape[0]):
             j = (i + 1) % polygon.shape[0]
             absolute_distance = self.compute_distance(xs, ys, polygon[i], polygon[j])

doctr/models/detection/predictor/pytorch.py

@@ -55,7 +55,8 @@ def forward(
             self.model, processed_batches, _params.device, _params.dtype
         )
         predicted_batches = [
-            self.model(batch, return_preds=True, return_model_output=True, **kwargs) for batch in processed_batches
+            self.model(batch, return_preds=True, return_model_output=return_maps, **kwargs)
+            for batch in processed_batches
         ]
         # Remove padding from loc predictions
         preds = _remove_padding(

doctr/models/layout/lw_detr/layers/pytorch.py

@@ -221,11 +221,12 @@ def forward(
                 .flatten(2)
                 .transpose(1, 2)
                 .reshape(batch_size * num_heads, hidden_dim, height, width)
+                .contiguous()
             )
             # batch_size, num_queries, num_heads, num_points, 2
             # -> batch_size, num_heads, num_queries, num_points, 2
             # -> batch_size*num_heads, num_queries, num_points, 2
-            sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1)
+            sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1).contiguous()
             # batch_size*num_heads, hidden_dim, num_queries, num_points
             sampling_value_l_ = nn.functional.grid_sample(
                 value_l_,

doctr/models/recognition/master/pytorch.py

@@ -249,6 +249,8 @@ def decode(self, encoded: torch.Tensor) -> torch.Tensor:
             output = self.decoder(ys, encoded, source_mask, target_mask)
             # update ys with the next token and ignore the first token (SOS)
             ys[:, i + 1] = self.linear(output[:, i]).argmax(-1)
+            if (ys == self.vocab_size).any(dim=-1).all():  # every sequence has emitted EOS
+                break
 
         # Shape (N, max_length, vocab_size + 1)
         return self.linear(output)

doctr/models/utils/pytorch.py

@@ -149,7 +149,9 @@ def set_device_and_dtype(
     Returns:
         the model and batches set
     """
-    model = model.to(device=device, dtype=dtype)
+    first = next(model.parameters(), None)
+    if first is None or first.device != torch.device(device) or first.dtype != dtype:
+        model = model.to(device=device, dtype=dtype)
     if isinstance(batches, tuple):
         return model, [
             (img.to(device=device, dtype=dtype), mask.to(device=device, dtype=torch.bool))