diff --git a/docs/source/using_doctr/using_model_export.rst b/docs/source/using_doctr/using_model_export.rst index 7cf94accf8..4aafce335f 100644 --- a/docs/source/using_doctr/using_model_export.rst +++ b/docs/source/using_doctr/using_model_export.rst @@ -11,19 +11,29 @@ 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 @@ -31,7 +41,7 @@ Advantages: reco_arch="crnn_mobilenet_v3_small", det_arch="linknet_resnet34", pretrained=True - ).cuda().half() + ).cuda().bfloat16() # or .half() for FP16 res = predictor(doc) diff --git a/docs/source/using_doctr/using_models.rst b/docs/source/using_doctr/using_models.rst index 0ba368e7e2..4775f66b31 100644 --- a/docs/source/using_doctr/using_models.rst +++ b/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? diff --git a/doctr/models/classification/predictor/pytorch.py b/doctr/models/classification/predictor/pytorch.py index 280743af4d..7c1627c212 100644 --- a/doctr/models/classification/predictor/pytorch.py +++ b/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] diff --git a/doctr/models/detection/differentiable_binarization/base.py b/doctr/models/detection/differentiable_binarization/base.py index a15fe771f6..0b21e12ba3 100644 --- a/doctr/models/detection/differentiable_binarization/base.py +++ b/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]) diff --git a/doctr/models/detection/predictor/pytorch.py b/doctr/models/detection/predictor/pytorch.py index b5ba3ec095..69d7e7bd03 100644 --- a/doctr/models/detection/predictor/pytorch.py +++ b/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( diff --git a/doctr/models/layout/lw_detr/layers/pytorch.py b/doctr/models/layout/lw_detr/layers/pytorch.py index 5456c71c9b..adc39bd844 100644 --- a/doctr/models/layout/lw_detr/layers/pytorch.py +++ b/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_, diff --git a/doctr/models/recognition/master/pytorch.py b/doctr/models/recognition/master/pytorch.py index e9fa84e3b8..0bbd2bf8fa 100644 --- a/doctr/models/recognition/master/pytorch.py +++ b/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) diff --git a/doctr/models/utils/pytorch.py b/doctr/models/utils/pytorch.py index 9fce34f20c..bc4b68b968 100644 --- a/doctr/models/utils/pytorch.py +++ b/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))