python-doctr 0.8.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,428 @@
1
+ # Copyright (C) 2021-2024, Mindee.
2
+
3
+ # This program is licensed under the Apache License 2.0.
4
+ # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
5
+
6
+ # Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
7
+
8
+ from copy import deepcopy
9
+ from typing import Any, Dict, List, Optional, Tuple, Union
10
+
11
+ import numpy as np
12
+ import tensorflow as tf
13
+ from tensorflow import keras
14
+ from tensorflow.keras import Sequential, layers
15
+
16
+ from doctr.file_utils import CLASS_NAME
17
+ from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, load_pretrained_params
18
+ from doctr.utils.repr import NestedObject
19
+
20
+ from ...classification import textnet_base, textnet_small, textnet_tiny
21
+ from ...modules.layers import FASTConvLayer
22
+ from .base import _FAST, FASTPostProcessor
23
+
24
+ __all__ = ["FAST", "fast_tiny", "fast_small", "fast_base", "reparameterize"]
25
+
26
+
27
+ default_cfgs: Dict[str, Dict[str, Any]] = {
28
+ "fast_tiny": {
29
+ "input_shape": (1024, 1024, 3),
30
+ "mean": (0.798, 0.785, 0.772),
31
+ "std": (0.264, 0.2749, 0.287),
32
+ "url": None,
33
+ },
34
+ "fast_small": {
35
+ "input_shape": (1024, 1024, 3),
36
+ "mean": (0.798, 0.785, 0.772),
37
+ "std": (0.264, 0.2749, 0.287),
38
+ "url": None,
39
+ },
40
+ "fast_base": {
41
+ "input_shape": (1024, 1024, 3),
42
+ "mean": (0.798, 0.785, 0.772),
43
+ "std": (0.264, 0.2749, 0.287),
44
+ "url": None,
45
+ },
46
+ }
47
+
48
+
49
+ class FastNeck(layers.Layer, NestedObject):
50
+ """Neck of the FAST architecture, composed of a series of 3x3 convolutions and upsampling layer.
51
+
52
+ Args:
53
+ ----
54
+ in_channels: number of input channels
55
+ out_channels: number of output channels
56
+ """
57
+
58
+ def __init__(
59
+ self,
60
+ in_channels: int,
61
+ out_channels: int = 128,
62
+ ) -> None:
63
+ super().__init__()
64
+ self.reduction = [FASTConvLayer(in_channels * scale, out_channels, kernel_size=3) for scale in [1, 2, 4, 8]]
65
+
66
+ def _upsample(self, x: tf.Tensor, y: tf.Tensor) -> tf.Tensor:
67
+ return tf.image.resize(x, size=y.shape[1:3], method="bilinear")
68
+
69
+ def call(self, x: tf.Tensor, **kwargs: Any) -> tf.Tensor:
70
+ f1, f2, f3, f4 = x
71
+ f1, f2, f3, f4 = [reduction(f, **kwargs) for reduction, f in zip(self.reduction, (f1, f2, f3, f4))]
72
+ f2, f3, f4 = [self._upsample(f, f1) for f in (f2, f3, f4)]
73
+ f = tf.concat((f1, f2, f3, f4), axis=-1)
74
+ return f
75
+
76
+
77
+ class FastHead(Sequential):
78
+ """Head of the FAST architecture
79
+
80
+ Args:
81
+ ----
82
+ in_channels: number of input channels
83
+ num_classes: number of output classes
84
+ out_channels: number of output channels
85
+ dropout: dropout probability
86
+ """
87
+
88
+ def __init__(
89
+ self,
90
+ in_channels: int,
91
+ num_classes: int,
92
+ out_channels: int = 128,
93
+ dropout: float = 0.1,
94
+ ) -> None:
95
+ _layers = [
96
+ FASTConvLayer(in_channels, out_channels, kernel_size=3),
97
+ layers.Dropout(dropout),
98
+ layers.Conv2D(num_classes, kernel_size=1, use_bias=False),
99
+ ]
100
+ super().__init__(_layers)
101
+
102
+
103
+ class FAST(_FAST, keras.Model, NestedObject):
104
+ """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
105
+ <https://arxiv.org/pdf/2111.02394.pdf>`_.
106
+
107
+ Args:
108
+ ----
109
+ feature extractor: the backbone serving as feature extractor
110
+ bin_thresh: threshold for binarization
111
+ box_thresh: minimal objectness score to consider a box
112
+ dropout_prob: dropout probability
113
+ pooling_size: size of the pooling layer
114
+ assume_straight_pages: if True, fit straight bounding boxes only
115
+ exportable: onnx exportable returns only logits
116
+ cfg: the configuration dict of the model
117
+ class_names: list of class names
118
+ """
119
+
120
+ _children_names: List[str] = ["feat_extractor", "neck", "head", "postprocessor"]
121
+
122
+ def __init__(
123
+ self,
124
+ feature_extractor: IntermediateLayerGetter,
125
+ bin_thresh: float = 0.3,
126
+ box_thresh: float = 0.1,
127
+ dropout_prob: float = 0.1,
128
+ pooling_size: int = 4, # different from paper performs better on close text-rich images
129
+ assume_straight_pages: bool = True,
130
+ exportable: bool = False,
131
+ cfg: Optional[Dict[str, Any]] = {},
132
+ class_names: List[str] = [CLASS_NAME],
133
+ ) -> None:
134
+ super().__init__()
135
+ self.class_names = class_names
136
+ num_classes: int = len(self.class_names)
137
+ self.cfg = cfg
138
+
139
+ self.feat_extractor = feature_extractor
140
+ self.exportable = exportable
141
+ self.assume_straight_pages = assume_straight_pages
142
+
143
+ # Identify the number of channels for the neck & head initialization
144
+ feat_out_channels = [
145
+ layers.Input(shape=in_shape[1:]).shape[-1] for in_shape in self.feat_extractor.output_shape
146
+ ]
147
+ # Initialize neck & head
148
+ self.neck = FastNeck(feat_out_channels[0], feat_out_channels[1])
149
+ self.head = FastHead(feat_out_channels[-1], num_classes, feat_out_channels[1], dropout_prob)
150
+
151
+ # NOTE: The post processing from the paper works not well for text-rich images
152
+ # so we use a modified version from DBNet
153
+ self.postprocessor = FASTPostProcessor(
154
+ assume_straight_pages=assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh
155
+ )
156
+
157
+ # Pooling layer as erosion reversal as described in the paper
158
+ self.pooling = layers.MaxPooling2D(pool_size=pooling_size // 2 + 1, strides=1, padding="same")
159
+
160
+ def compute_loss(
161
+ self,
162
+ out_map: tf.Tensor,
163
+ target: List[Dict[str, np.ndarray]],
164
+ eps: float = 1e-6,
165
+ ) -> tf.Tensor:
166
+ """Compute fast loss, 2 x Dice loss where the text kernel loss is scaled by 0.5.
167
+
168
+ Args:
169
+ ----
170
+ out_map: output feature map of the model of shape (N, num_classes, H, W)
171
+ target: list of dictionary where each dict has a `boxes` and a `flags` entry
172
+ eps: epsilon factor in dice loss
173
+
174
+ Returns:
175
+ -------
176
+ A loss tensor
177
+ """
178
+ targets = self.build_target(target, out_map.shape[1:], True)
179
+
180
+ seg_target = tf.convert_to_tensor(targets[0], dtype=out_map.dtype)
181
+ seg_mask = tf.convert_to_tensor(targets[1], dtype=out_map.dtype)
182
+ shrunken_kernel = tf.convert_to_tensor(targets[2], dtype=out_map.dtype)
183
+
184
+ def ohem(score: tf.Tensor, gt: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
185
+ pos_num = tf.reduce_sum(tf.cast(gt > 0.5, dtype=tf.int32)) - tf.reduce_sum(
186
+ tf.cast((gt > 0.5) & (mask <= 0.5), dtype=tf.int32)
187
+ )
188
+ neg_num = tf.reduce_sum(tf.cast(gt <= 0.5, dtype=tf.int32))
189
+ neg_num = tf.minimum(pos_num * 3, neg_num)
190
+
191
+ if neg_num == 0 or pos_num == 0:
192
+ return mask
193
+
194
+ neg_score_sorted, _ = tf.nn.top_k(-tf.boolean_mask(score, gt <= 0.5), k=neg_num)
195
+ threshold = -neg_score_sorted[-1]
196
+
197
+ selected_mask = tf.math.logical_and((score >= threshold) | (gt > 0.5), (mask > 0.5))
198
+ return tf.cast(selected_mask, dtype=tf.float32)
199
+
200
+ if len(self.class_names) > 1:
201
+ kernels = tf.nn.softmax(out_map, axis=-1)
202
+ prob_map = tf.nn.softmax(self.pooling(out_map), axis=-1)
203
+ else:
204
+ kernels = tf.sigmoid(out_map)
205
+ prob_map = tf.sigmoid(self.pooling(out_map))
206
+
207
+ # As described in the paper, we use the Dice loss for the text segmentation map and the Dice loss scaled by 0.5.
208
+ selected_masks = tf.stack(
209
+ [ohem(score, gt, mask) for score, gt, mask in zip(prob_map, seg_target, seg_mask)], axis=0
210
+ )
211
+ inter = tf.reduce_sum(selected_masks * prob_map * seg_target, axis=(0, 1, 2))
212
+ cardinality = tf.reduce_sum(selected_masks * (prob_map + seg_target), axis=(0, 1, 2))
213
+ text_loss = tf.reduce_mean((1 - 2 * inter / (cardinality + eps))) * 0.5
214
+
215
+ # As described in the paper, we use the Dice loss for the text kernel map.
216
+ selected_masks = seg_target * seg_mask
217
+ inter = tf.reduce_sum(selected_masks * kernels * shrunken_kernel, axis=(0, 1, 2))
218
+ cardinality = tf.reduce_sum(selected_masks * (kernels + shrunken_kernel), axis=(0, 1, 2))
219
+ kernel_loss = tf.reduce_mean((1 - 2 * inter / (cardinality + eps)))
220
+
221
+ return text_loss + kernel_loss
222
+
223
+ def call(
224
+ self,
225
+ x: tf.Tensor,
226
+ target: Optional[List[Dict[str, np.ndarray]]] = None,
227
+ return_model_output: bool = False,
228
+ return_preds: bool = False,
229
+ **kwargs: Any,
230
+ ) -> Dict[str, Any]:
231
+ feat_maps = self.feat_extractor(x, **kwargs)
232
+ # Pass through the Neck & Head & Upsample
233
+ feat_concat = self.neck(feat_maps, **kwargs)
234
+ logits: tf.Tensor = self.head(feat_concat, **kwargs)
235
+ logits = layers.UpSampling2D(size=x.shape[-2] // logits.shape[-2], interpolation="bilinear")(logits, **kwargs)
236
+
237
+ out: Dict[str, tf.Tensor] = {}
238
+ if self.exportable:
239
+ out["logits"] = logits
240
+ return out
241
+
242
+ if return_model_output or target is None or return_preds:
243
+ prob_map = _bf16_to_float32(tf.math.sigmoid(self.pooling(logits, **kwargs)))
244
+
245
+ if return_model_output:
246
+ out["out_map"] = prob_map
247
+
248
+ if target is None or return_preds:
249
+ # Post-process boxes (keep only text predictions)
250
+ out["preds"] = [dict(zip(self.class_names, preds)) for preds in self.postprocessor(prob_map.numpy())]
251
+
252
+ if target is not None:
253
+ loss = self.compute_loss(logits, target)
254
+ out["loss"] = loss
255
+
256
+ return out
257
+
258
+
259
+ def reparameterize(model: Union[FAST, layers.Layer]) -> FAST:
260
+ """Fuse batchnorm and conv layers and reparameterize the model
261
+
262
+ args:
263
+ ----
264
+ model: the FAST model to reparameterize
265
+
266
+ Returns:
267
+ -------
268
+ the reparameterized model
269
+ """
270
+ last_conv = None
271
+ last_conv_idx = None
272
+
273
+ for idx, layer in enumerate(model.layers):
274
+ if hasattr(layer, "layers") or isinstance(
275
+ layer, (FASTConvLayer, FastNeck, FastHead, layers.BatchNormalization, layers.Conv2D)
276
+ ):
277
+ if isinstance(layer, layers.BatchNormalization):
278
+ # fuse batchnorm only if it is followed by a conv layer
279
+ if last_conv is None:
280
+ continue
281
+ conv_w = last_conv.kernel
282
+ conv_b = last_conv.bias if last_conv.use_bias else tf.zeros_like(layer.moving_mean)
283
+
284
+ factor = layer.gamma / tf.sqrt(layer.moving_variance + layer.epsilon)
285
+ last_conv.kernel = conv_w * factor.numpy().reshape([1, 1, 1, -1])
286
+ if last_conv.use_bias:
287
+ last_conv.bias.assign((conv_b - layer.moving_mean) * factor + layer.beta)
288
+ model.layers[last_conv_idx] = last_conv # Replace the last conv layer with the fused version
289
+ model.layers[idx] = layers.Lambda(lambda x: x)
290
+ last_conv = None
291
+ elif isinstance(layer, layers.Conv2D):
292
+ last_conv = layer
293
+ last_conv_idx = idx
294
+ elif isinstance(layer, FASTConvLayer):
295
+ layer.reparameterize_layer()
296
+ elif isinstance(layer, FastNeck):
297
+ for reduction in layer.reduction:
298
+ reduction.reparameterize_layer()
299
+ elif isinstance(layer, FastHead):
300
+ reparameterize(layer)
301
+ else:
302
+ reparameterize(layer)
303
+ return model
304
+
305
+
306
+ def _fast(
307
+ arch: str,
308
+ pretrained: bool,
309
+ backbone_fn,
310
+ feat_layers: List[str],
311
+ pretrained_backbone: bool = True,
312
+ input_shape: Optional[Tuple[int, int, int]] = None,
313
+ **kwargs: Any,
314
+ ) -> FAST:
315
+ pretrained_backbone = pretrained_backbone and not pretrained
316
+
317
+ # Patch the config
318
+ _cfg = deepcopy(default_cfgs[arch])
319
+ _cfg["input_shape"] = input_shape or _cfg["input_shape"]
320
+ if not kwargs.get("class_names", None):
321
+ kwargs["class_names"] = _cfg.get("class_names", [CLASS_NAME])
322
+ else:
323
+ kwargs["class_names"] = sorted(kwargs["class_names"])
324
+
325
+ # Feature extractor
326
+ feat_extractor = IntermediateLayerGetter(
327
+ backbone_fn(
328
+ input_shape=_cfg["input_shape"],
329
+ include_top=False,
330
+ pretrained=pretrained_backbone,
331
+ ),
332
+ feat_layers,
333
+ )
334
+
335
+ # Build the model
336
+ model = FAST(feat_extractor, cfg=_cfg, **kwargs)
337
+ # Load pretrained parameters
338
+ if pretrained:
339
+ load_pretrained_params(model, _cfg["url"])
340
+
341
+ # Build the model for reparameterization to access the layers
342
+ _ = model(tf.random.uniform(shape=[1, *_cfg["input_shape"]], maxval=1, dtype=tf.float32), training=False)
343
+
344
+ return model
345
+
346
+
347
+ def fast_tiny(pretrained: bool = False, **kwargs: Any) -> FAST:
348
+ """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
349
+ <https://arxiv.org/pdf/2111.02394.pdf>`_, using a tiny TextNet backbone.
350
+
351
+ >>> import tensorflow as tf
352
+ >>> from doctr.models import fast_tiny
353
+ >>> model = fast_tiny(pretrained=True)
354
+ >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
355
+ >>> out = model(input_tensor)
356
+
357
+ Args:
358
+ ----
359
+ pretrained (bool): If True, returns a model pre-trained on our text detection dataset
360
+ **kwargs: keyword arguments of the DBNet architecture
361
+
362
+ Returns:
363
+ -------
364
+ text detection architecture
365
+ """
366
+ return _fast(
367
+ "fast_tiny",
368
+ pretrained,
369
+ textnet_tiny,
370
+ ["stage_0", "stage_1", "stage_2", "stage_3"],
371
+ **kwargs,
372
+ )
373
+
374
+
375
+ def fast_small(pretrained: bool = False, **kwargs: Any) -> FAST:
376
+ """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
377
+ <https://arxiv.org/pdf/2111.02394.pdf>`_, using a small TextNet backbone.
378
+
379
+ >>> import tensorflow as tf
380
+ >>> from doctr.models import fast_small
381
+ >>> model = fast_small(pretrained=True)
382
+ >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
383
+ >>> out = model(input_tensor)
384
+
385
+ Args:
386
+ ----
387
+ pretrained (bool): If True, returns a model pre-trained on our text detection dataset
388
+ **kwargs: keyword arguments of the DBNet architecture
389
+
390
+ Returns:
391
+ -------
392
+ text detection architecture
393
+ """
394
+ return _fast(
395
+ "fast_small",
396
+ pretrained,
397
+ textnet_small,
398
+ ["stage_0", "stage_1", "stage_2", "stage_3"],
399
+ **kwargs,
400
+ )
401
+
402
+
403
+ def fast_base(pretrained: bool = False, **kwargs: Any) -> FAST:
404
+ """FAST as described in `"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation"
405
+ <https://arxiv.org/pdf/2111.02394.pdf>`_, using a base TextNet backbone.
406
+
407
+ >>> import tensorflow as tf
408
+ >>> from doctr.models import fast_base
409
+ >>> model = fast_base(pretrained=True)
410
+ >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
411
+ >>> out = model(input_tensor)
412
+
413
+ Args:
414
+ ----
415
+ pretrained (bool): If True, returns a model pre-trained on our text detection dataset
416
+ **kwargs: keyword arguments of the DBNet architecture
417
+
418
+ Returns:
419
+ -------
420
+ text detection architecture
421
+ """
422
+ return _fast(
423
+ "fast_base",
424
+ pretrained,
425
+ textnet_base,
426
+ ["stage_0", "stage_1", "stage_2", "stage_3"],
427
+ **kwargs,
428
+ )
@@ -17,7 +17,16 @@ ARCHS: List[str]
17
17
 
18
18
 
19
19
  if is_tf_available():
20
- ARCHS = ["db_resnet50", "db_mobilenet_v3_large", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
20
+ ARCHS = [
21
+ "db_resnet50",
22
+ "db_mobilenet_v3_large",
23
+ "linknet_resnet18",
24
+ "linknet_resnet34",
25
+ "linknet_resnet50",
26
+ "fast_tiny",
27
+ "fast_small",
28
+ "fast_base",
29
+ ]
21
30
  elif is_torch_available():
22
31
  ARCHS = [
23
32
  "db_resnet34",
@@ -26,6 +35,9 @@ elif is_torch_available():
26
35
  "linknet_resnet18",
27
36
  "linknet_resnet34",
28
37
  "linknet_resnet50",
38
+ "fast_tiny",
39
+ "fast_small",
40
+ "fast_base",
29
41
  ]
30
42
 
31
43
 
@@ -40,7 +52,7 @@ def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True,
40
52
  assume_straight_pages=assume_straight_pages,
41
53
  )
42
54
  else:
43
- if not isinstance(arch, (detection.DBNet, detection.LinkNet)):
55
+ if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):
44
56
  raise ValueError(f"unknown architecture: {type(arch)}")
45
57
 
46
58
  _model = arch
@@ -5,6 +5,7 @@
5
5
 
6
6
  from typing import Tuple, Union
7
7
 
8
+ import numpy as np
8
9
  import torch
9
10
  import torch.nn as nn
10
11
 
@@ -26,18 +27,20 @@ class FASTConvLayer(nn.Module):
26
27
  ) -> None:
27
28
  super().__init__()
28
29
 
29
- converted_ks = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
30
+ self.groups = groups
31
+ self.in_channels = in_channels
32
+ self.converted_ks = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
30
33
 
31
34
  self.hor_conv, self.hor_bn = None, None
32
35
  self.ver_conv, self.ver_bn = None, None
33
36
 
34
- padding = (int(((converted_ks[0] - 1) * dilation) / 2), int(((converted_ks[1] - 1) * dilation) / 2))
37
+ padding = (int(((self.converted_ks[0] - 1) * dilation) / 2), int(((self.converted_ks[1] - 1) * dilation) / 2))
35
38
 
36
39
  self.activation = nn.ReLU(inplace=True)
37
40
  self.conv = nn.Conv2d(
38
41
  in_channels,
39
42
  out_channels,
40
- kernel_size=converted_ks,
43
+ kernel_size=self.converted_ks,
41
44
  stride=stride,
42
45
  padding=padding,
43
46
  dilation=dilation,
@@ -47,12 +50,12 @@ class FASTConvLayer(nn.Module):
47
50
 
48
51
  self.bn = nn.BatchNorm2d(out_channels)
49
52
 
50
- if converted_ks[1] != 1:
53
+ if self.converted_ks[1] != 1:
51
54
  self.ver_conv = nn.Conv2d(
52
55
  in_channels,
53
56
  out_channels,
54
- kernel_size=(converted_ks[0], 1),
55
- padding=(int(((converted_ks[0] - 1) * dilation) / 2), 0),
57
+ kernel_size=(self.converted_ks[0], 1),
58
+ padding=(int(((self.converted_ks[0] - 1) * dilation) / 2), 0),
56
59
  stride=stride,
57
60
  dilation=dilation,
58
61
  groups=groups,
@@ -60,12 +63,12 @@ class FASTConvLayer(nn.Module):
60
63
  )
61
64
  self.ver_bn = nn.BatchNorm2d(out_channels)
62
65
 
63
- if converted_ks[0] != 1:
66
+ if self.converted_ks[0] != 1:
64
67
  self.hor_conv = nn.Conv2d(
65
68
  in_channels,
66
69
  out_channels,
67
- kernel_size=(1, converted_ks[1]),
68
- padding=(0, int(((converted_ks[1] - 1) * dilation) / 2)),
70
+ kernel_size=(1, self.converted_ks[1]),
71
+ padding=(0, int(((self.converted_ks[1] - 1) * dilation) / 2)),
69
72
  stride=stride,
70
73
  dilation=dilation,
71
74
  groups=groups,
@@ -76,6 +79,9 @@ class FASTConvLayer(nn.Module):
76
79
  self.rbr_identity = nn.BatchNorm2d(in_channels) if out_channels == in_channels and stride == 1 else None
77
80
 
78
81
  def forward(self, x: torch.Tensor) -> torch.Tensor:
82
+ if hasattr(self, "fused_conv"):
83
+ return self.activation(self.fused_conv(x))
84
+
79
85
  main_outputs = self.bn(self.conv(x))
80
86
  vertical_outputs = self.ver_bn(self.ver_conv(x)) if self.ver_conv is not None and self.ver_bn is not None else 0
81
87
  horizontal_outputs = (
@@ -84,3 +90,77 @@ class FASTConvLayer(nn.Module):
84
90
  id_out = self.rbr_identity(x) if self.rbr_identity is not None and self.ver_bn is not None else 0
85
91
 
86
92
  return self.activation(main_outputs + vertical_outputs + horizontal_outputs + id_out)
93
+
94
+ # The following logic is used to reparametrize the layer
95
+ # Borrowed from: https://github.com/czczup/FAST/blob/main/models/utils/nas_utils.py
96
+ def _identity_to_conv(
97
+ self, identity: Union[nn.BatchNorm2d, None]
98
+ ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]:
99
+ if identity is None or identity.running_var is None:
100
+ return 0, 0
101
+ if not hasattr(self, "id_tensor"):
102
+ input_dim = self.in_channels // self.groups
103
+ kernel_value = np.zeros((self.in_channels, input_dim, 1, 1), dtype=np.float32)
104
+ for i in range(self.in_channels):
105
+ kernel_value[i, i % input_dim, 0, 0] = 1
106
+ id_tensor = torch.from_numpy(kernel_value).to(identity.weight.device)
107
+ self.id_tensor = self._pad_to_mxn_tensor(id_tensor)
108
+ kernel = self.id_tensor
109
+ std = (identity.running_var + identity.eps).sqrt() # type: ignore[attr-defined]
110
+ t = (identity.weight / std).reshape(-1, 1, 1, 1)
111
+ return kernel * t, identity.bias - identity.running_mean * identity.weight / std
112
+
113
+ def _fuse_bn_tensor(self, conv: nn.Conv2d, bn: nn.BatchNorm2d) -> Tuple[torch.Tensor, torch.Tensor]:
114
+ kernel = conv.weight
115
+ kernel = self._pad_to_mxn_tensor(kernel)
116
+ std = (bn.running_var + bn.eps).sqrt() # type: ignore
117
+ t = (bn.weight / std).reshape(-1, 1, 1, 1)
118
+ return kernel * t, bn.bias - bn.running_mean * bn.weight / std
119
+
120
+ def _get_equivalent_kernel_bias(self) -> Tuple[torch.Tensor, torch.Tensor]:
121
+ kernel_mxn, bias_mxn = self._fuse_bn_tensor(self.conv, self.bn)
122
+ if self.ver_conv is not None:
123
+ kernel_mx1, bias_mx1 = self._fuse_bn_tensor(self.ver_conv, self.ver_bn) # type: ignore[arg-type]
124
+ else:
125
+ kernel_mx1, bias_mx1 = 0, 0 # type: ignore[assignment]
126
+ if self.hor_conv is not None:
127
+ kernel_1xn, bias_1xn = self._fuse_bn_tensor(self.hor_conv, self.hor_bn) # type: ignore[arg-type]
128
+ else:
129
+ kernel_1xn, bias_1xn = 0, 0 # type: ignore[assignment]
130
+ kernel_id, bias_id = self._identity_to_conv(self.rbr_identity)
131
+ kernel_mxn = kernel_mxn + kernel_mx1 + kernel_1xn + kernel_id
132
+ bias_mxn = bias_mxn + bias_mx1 + bias_1xn + bias_id
133
+ return kernel_mxn, bias_mxn
134
+
135
+ def _pad_to_mxn_tensor(self, kernel: torch.Tensor) -> torch.Tensor:
136
+ kernel_height, kernel_width = self.converted_ks
137
+ height, width = kernel.shape[2:]
138
+ pad_left_right = (kernel_width - width) // 2
139
+ pad_top_down = (kernel_height - height) // 2
140
+ return torch.nn.functional.pad(kernel, [pad_left_right, pad_left_right, pad_top_down, pad_top_down], value=0)
141
+
142
+ def reparameterize_layer(self):
143
+ if hasattr(self, "fused_conv"):
144
+ return
145
+ kernel, bias = self._get_equivalent_kernel_bias()
146
+ self.fused_conv = nn.Conv2d(
147
+ in_channels=self.conv.in_channels,
148
+ out_channels=self.conv.out_channels,
149
+ kernel_size=self.conv.kernel_size, # type: ignore[arg-type]
150
+ stride=self.conv.stride, # type: ignore[arg-type]
151
+ padding=self.conv.padding, # type: ignore[arg-type]
152
+ dilation=self.conv.dilation, # type: ignore[arg-type]
153
+ groups=self.conv.groups,
154
+ bias=True,
155
+ )
156
+ self.fused_conv.weight.data = kernel
157
+ self.fused_conv.bias.data = bias # type: ignore[union-attr]
158
+ self.deploy = True
159
+ for para in self.parameters():
160
+ para.detach_()
161
+ for attr in ["conv", "bn", "ver_conv", "ver_bn", "hor_conv", "hor_bn"]:
162
+ if hasattr(self, attr):
163
+ self.__delattr__(attr)
164
+
165
+ if hasattr(self, "rbr_identity"):
166
+ self.__delattr__("rbr_identity")