We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Batch inference during evaluation in text detection training
I'm training a `DBNet_r50dcn` model with `tools/train.py` script and, even though training epoch is executed in batch with `samples_per_gpu`, during evaluation one image at time is fed to the model. Clearly this makes the whole execution extremely slow if we have a large validation set.

Is there a solution to run in batch mode also the evaluation step?

To summarize: same problem of [this comment](https://github.com/open-mmlab/mmocr/issues/294#issuecomment-864533214), but in my case is for text detection.



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mmocr/models/textdet/losses/db_loss.py`
Content:
```
1 import torch
2 import torch.nn.functional as F
3 from mmdet.models.builder import LOSSES
4 from torch import nn
5 
6 from mmocr.models.common.losses.dice_loss import DiceLoss
7 
8 
9 @LOSSES.register_module()
10 class DBLoss(nn.Module):
11     """The class for implementing DBNet loss.
12 
13     This is partially adapted from https://github.com/MhLiao/DB.
14     """
15 
16     def __init__(self,
17                  alpha=1,
18                  beta=1,
19                  reduction='mean',
20                  negative_ratio=3.0,
21                  eps=1e-6,
22                  bbce_loss=False):
23         """Initialization.
24 
25         Args:
26             alpha (float): The binary loss coef.
27             beta (float): The threshold loss coef.
28             reduction (str): The way to reduce the loss.
29             negative_ratio (float): The ratio of positives to negatives.
30             eps (float): Epsilon in the threshold loss function.
31             bbce_loss (bool): Whether to use balanced bce for probability loss.
32                 If False, dice loss will be used instead.
33         """
34         super().__init__()
35         assert reduction in ['mean',
36                              'sum'], " reduction must in ['mean','sum']"
37         self.alpha = alpha
38         self.beta = beta
39         self.reduction = reduction
40         self.negative_ratio = negative_ratio
41         self.eps = eps
42         self.bbce_loss = bbce_loss
43         self.dice_loss = DiceLoss(eps=eps)
44 
45     def bitmasks2tensor(self, bitmasks, target_sz):
46         """Convert Bitmasks to tensor.
47 
48         Args:
49             bitmasks (list[BitMasks]): The BitMasks list. Each item is for
50                 one img.
51             target_sz (tuple(int, int)): The target tensor size of KxHxW
52                 with K being the number of kernels.
53 
54         Returns
55             result_tensors (list[tensor]): The list of kernel tensors. Each
56                 element is for one kernel level.
57         """
58         assert isinstance(bitmasks, list)
59         assert isinstance(target_sz, tuple)
60 
61         batch_size = len(bitmasks)
62         num_levels = len(bitmasks[0])
63 
64         result_tensors = []
65 
66         for level_inx in range(num_levels):
67             kernel = []
68             for batch_inx in range(batch_size):
69                 mask = torch.from_numpy(bitmasks[batch_inx].masks[level_inx])
70                 mask_sz = mask.shape
71                 pad = [
72                     0, target_sz[1] - mask_sz[1], 0, target_sz[0] - mask_sz[0]
73                 ]
74                 mask = F.pad(mask, pad, mode='constant', value=0)
75                 kernel.append(mask)
76             kernel = torch.stack(kernel)
77             result_tensors.append(kernel)
78 
79         return result_tensors
80 
81     def balance_bce_loss(self, pred, gt, mask):
82 
83         positive = (gt * mask)
84         negative = ((1 - gt) * mask)
85         positive_count = int(positive.float().sum())
86         negative_count = min(
87             int(negative.float().sum()),
88             int(positive_count * self.negative_ratio))
89 
90         assert gt.max() <= 1 and gt.min() >= 0
91         assert pred.max() <= 1 and pred.min() >= 0
92         loss = F.binary_cross_entropy(pred, gt, reduction='none')
93         positive_loss = loss * positive.float()
94         negative_loss = loss * negative.float()
95 
96         negative_loss, _ = torch.topk(negative_loss.view(-1), negative_count)
97 
98         balance_loss = (positive_loss.sum() + negative_loss.sum()) / (
99             positive_count + negative_count + self.eps)
100 
101         return balance_loss
102 
103     def l1_thr_loss(self, pred, gt, mask):
104         thr_loss = torch.abs((pred - gt) * mask).sum() / (
105             mask.sum() + self.eps)
106         return thr_loss
107 
108     def forward(self, preds, downsample_ratio, gt_shrink, gt_shrink_mask,
109                 gt_thr, gt_thr_mask):
110         """Compute DBNet loss.
111 
112         Args:
113             preds (tensor): The output tensor with size of Nx3xHxW.
114             downsample_ratio (float): The downsample ratio for the
115                 ground truths.
116             gt_shrink (list[BitmapMasks]): The mask list with each element
117                 being the shrinked text mask for one img.
118             gt_shrink_mask (list[BitmapMasks]): The effective mask list with
119                 each element being the shrinked effective mask for one img.
120             gt_thr (list[BitmapMasks]): The mask list with each element
121                 being the threshold text mask for one img.
122             gt_thr_mask (list[BitmapMasks]): The effective mask list with
123                 each element being the threshold effective mask for one img.
124 
125         Returns:
126             results(dict): The dict for dbnet losses with loss_prob,
127                 loss_db and loss_thresh.
128         """
129         assert isinstance(downsample_ratio, float)
130 
131         assert isinstance(gt_shrink, list)
132         assert isinstance(gt_shrink_mask, list)
133         assert isinstance(gt_thr, list)
134         assert isinstance(gt_thr_mask, list)
135 
136         preds = preds[0]
137 
138         pred_prob = preds[:, 0, :, :]
139         pred_thr = preds[:, 1, :, :]
140         pred_db = preds[:, 2, :, :]
141         feature_sz = preds.size()
142 
143         keys = ['gt_shrink', 'gt_shrink_mask', 'gt_thr', 'gt_thr_mask']
144         gt = {}
145         for k in keys:
146             gt[k] = eval(k)
147             gt[k] = [item.rescale(downsample_ratio) for item in gt[k]]
148             gt[k] = self.bitmasks2tensor(gt[k], feature_sz[2:])
149             gt[k] = [item.to(preds.device) for item in gt[k]]
150         gt['gt_shrink'][0] = (gt['gt_shrink'][0] > 0).float()
151         if self.bbce_loss:
152             loss_prob = self.balance_bce_loss(pred_prob, gt['gt_shrink'][0],
153                                               gt['gt_shrink_mask'][0])
154         else:
155             loss_prob = self.dice_loss(pred_prob, gt['gt_shrink'][0],
156                                        gt['gt_shrink_mask'][0])
157 
158         loss_db = self.dice_loss(pred_db, gt['gt_shrink'][0],
159                                  gt['gt_shrink_mask'][0])
160         loss_thr = self.l1_thr_loss(pred_thr, gt['gt_thr'][0],
161                                     gt['gt_thr_mask'][0])
162 
163         results = dict(
164             loss_prob=self.alpha * loss_prob,
165             loss_db=loss_db,
166             loss_thr=self.beta * loss_thr)
167 
168         return results
169
```

Path: `mmocr/models/textdet/dense_heads/db_head.py`
Content:
```
1 import torch
2 import torch.nn as nn
3 from mmcv.runner import BaseModule, Sequential
4 from mmdet.models.builder import HEADS, build_loss
5 
6 from .head_mixin import HeadMixin
7 
8 
9 @HEADS.register_module()
10 class DBHead(HeadMixin, BaseModule):
11     """The class for DBNet head.
12 
13     This was partially adapted from https://github.com/MhLiao/DB
14     """
15 
16     def __init__(self,
17                  in_channels,
18                  with_bias=False,
19                  decoding_type='db',
20                  text_repr_type='poly',
21                  downsample_ratio=1.0,
22                  loss=dict(type='DBLoss'),
23                  train_cfg=None,
24                  test_cfg=None,
25                  init_cfg=[
26                      dict(type='Kaiming', layer='Conv'),
27                      dict(
28                          type='Constant', layer='BatchNorm', val=1., bias=1e-4)
29                  ]):
30         """Initialization.
31 
32         Args:
33             in_channels (int): The number of input channels of the db head.
34             decoding_type (str): The type of decoder for dbnet.
35             text_repr_type (str): Boundary encoding type 'poly' or 'quad'.
36             downsample_ratio (float): The downsample ratio of ground truths.
37             loss (dict): The type of loss for dbnet.
38         """
39         super().__init__(init_cfg=init_cfg)
40 
41         assert isinstance(in_channels, int)
42 
43         self.in_channels = in_channels
44         self.text_repr_type = text_repr_type
45         self.loss_module = build_loss(loss)
46         self.train_cfg = train_cfg
47         self.test_cfg = test_cfg
48         self.downsample_ratio = downsample_ratio
49         self.decoding_type = decoding_type
50 
51         self.binarize = Sequential(
52             nn.Conv2d(
53                 in_channels, in_channels // 4, 3, bias=with_bias, padding=1),
54             nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplace=True),
55             nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 2, 2),
56             nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplace=True),
57             nn.ConvTranspose2d(in_channels // 4, 1, 2, 2), nn.Sigmoid())
58 
59         self.threshold = self._init_thr(in_channels)
60 
61     def diff_binarize(self, prob_map, thr_map, k):
62         return torch.reciprocal(1.0 + torch.exp(-k * (prob_map - thr_map)))
63 
64     def forward(self, inputs):
65         prob_map = self.binarize(inputs)
66         thr_map = self.threshold(inputs)
67         binary_map = self.diff_binarize(prob_map, thr_map, k=50)
68         outputs = torch.cat((prob_map, thr_map, binary_map), dim=1)
69         return (outputs, )
70 
71     def _init_thr(self, inner_channels, bias=False):
72         in_channels = inner_channels
73         seq = Sequential(
74             nn.Conv2d(
75                 in_channels, inner_channels // 4, 3, padding=1, bias=bias),
76             nn.BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),
77             nn.ConvTranspose2d(inner_channels // 4, inner_channels // 4, 2, 2),
78             nn.BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),
79             nn.ConvTranspose2d(inner_channels // 4, 1, 2, 2), nn.Sigmoid())
80         return seq
81
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mmocr/models/textdet/dense_heads/db_head.py b/mmocr/models/textdet/dense_heads/db_head.py
--- a/mmocr/models/textdet/dense_heads/db_head.py
+++ b/mmocr/models/textdet/dense_heads/db_head.py
@@ -66,7 +66,7 @@
         thr_map = self.threshold(inputs)
         binary_map = self.diff_binarize(prob_map, thr_map, k=50)
         outputs = torch.cat((prob_map, thr_map, binary_map), dim=1)
-        return (outputs, )
+        return outputs
 
     def _init_thr(self, inner_channels, bias=False):
         in_channels = inner_channels
diff --git a/mmocr/models/textdet/losses/db_loss.py b/mmocr/models/textdet/losses/db_loss.py
--- a/mmocr/models/textdet/losses/db_loss.py
+++ b/mmocr/models/textdet/losses/db_loss.py
@@ -133,8 +133,6 @@
         assert isinstance(gt_thr, list)
         assert isinstance(gt_thr_mask, list)
 
-        preds = preds[0]
-
         pred_prob = preds[:, 0, :, :]
         pred_thr = preds[:, 1, :, :]
         pred_db = preds[:, 2, :, :]

{"golden_diff": "diff --git a/mmocr/models/textdet/dense_heads/db_head.py b/mmocr/models/textdet/dense_heads/db_head.py\n--- a/mmocr/models/textdet/dense_heads/db_head.py\n+++ b/mmocr/models/textdet/dense_heads/db_head.py\n@@ -66,7 +66,7 @@\n         thr_map = self.threshold(inputs)\n         binary_map = self.diff_binarize(prob_map, thr_map, k=50)\n         outputs = torch.cat((prob_map, thr_map, binary_map), dim=1)\n-        return (outputs, )\n+        return outputs\n \n     def _init_thr(self, inner_channels, bias=False):\n         in_channels = inner_channels\ndiff --git a/mmocr/models/textdet/losses/db_loss.py b/mmocr/models/textdet/losses/db_loss.py\n--- a/mmocr/models/textdet/losses/db_loss.py\n+++ b/mmocr/models/textdet/losses/db_loss.py\n@@ -133,8 +133,6 @@\n         assert isinstance(gt_thr, list)\n         assert isinstance(gt_thr_mask, list)\n \n-        preds = preds[0]\n-\n         pred_prob = preds[:, 0, :, :]\n         pred_thr = preds[:, 1, :, :]\n         pred_db = preds[:, 2, :, :]\n", "issue": "Batch inference during evaluation in text detection training\nI'm training a `DBNet_r50dcn` model with `tools/train.py` script and, even though training epoch is executed in batch\u00a0with `samples_per_gpu`, during evaluation one image at time is fed to the model. Clearly this makes the whole execution extremely slow if we have a large validation set.\r\n\r\nIs there a solution to run in batch mode also the evaluation step?\r\n\r\nTo summarize: same problem of [this comment](https://github.com/open-mmlab/mmocr/issues/294#issuecomment-864533214), but in my case is for text detection.\r\n\r\n\n", "before_files": [{"content": "import torch\nimport torch.nn.functional as F\nfrom mmdet.models.builder import LOSSES\nfrom torch import nn\n\nfrom mmocr.models.common.losses.dice_loss import DiceLoss\n\n\n@LOSSES.register_module()\nclass DBLoss(nn.Module):\n    \"\"\"The class for implementing DBNet loss.\n\n    This is partially adapted from https://github.com/MhLiao/DB.\n    \"\"\"\n\n    def __init__(self,\n                 alpha=1,\n                 beta=1,\n                 reduction='mean',\n                 negative_ratio=3.0,\n                 eps=1e-6,\n                 bbce_loss=False):\n        \"\"\"Initialization.\n\n        Args:\n            alpha (float): The binary loss coef.\n            beta (float): The threshold loss coef.\n            reduction (str): The way to reduce the loss.\n            negative_ratio (float): The ratio of positives to negatives.\n            eps (float): Epsilon in the threshold loss function.\n            bbce_loss (bool): Whether to use balanced bce for probability loss.\n                If False, dice loss will be used instead.\n        \"\"\"\n        super().__init__()\n        assert reduction in ['mean',\n                             'sum'], \" reduction must in ['mean','sum']\"\n        self.alpha = alpha\n        self.beta = beta\n        self.reduction = reduction\n        self.negative_ratio = negative_ratio\n        self.eps = eps\n        self.bbce_loss = bbce_loss\n        self.dice_loss = DiceLoss(eps=eps)\n\n    def bitmasks2tensor(self, bitmasks, target_sz):\n        \"\"\"Convert Bitmasks to tensor.\n\n        Args:\n            bitmasks (list[BitMasks]): The BitMasks list. Each item is for\n                one img.\n            target_sz (tuple(int, int)): The target tensor size of KxHxW\n                with K being the number of kernels.\n\n        Returns\n            result_tensors (list[tensor]): The list of kernel tensors. Each\n                element is for one kernel level.\n        \"\"\"\n        assert isinstance(bitmasks, list)\n        assert isinstance(target_sz, tuple)\n\n        batch_size = len(bitmasks)\n        num_levels = len(bitmasks[0])\n\n        result_tensors = []\n\n        for level_inx in range(num_levels):\n            kernel = []\n            for batch_inx in range(batch_size):\n                mask = torch.from_numpy(bitmasks[batch_inx].masks[level_inx])\n                mask_sz = mask.shape\n                pad = [\n                    0, target_sz[1] - mask_sz[1], 0, target_sz[0] - mask_sz[0]\n                ]\n                mask = F.pad(mask, pad, mode='constant', value=0)\n                kernel.append(mask)\n            kernel = torch.stack(kernel)\n            result_tensors.append(kernel)\n\n        return result_tensors\n\n    def balance_bce_loss(self, pred, gt, mask):\n\n        positive = (gt * mask)\n        negative = ((1 - gt) * mask)\n        positive_count = int(positive.float().sum())\n        negative_count = min(\n            int(negative.float().sum()),\n            int(positive_count * self.negative_ratio))\n\n        assert gt.max() <= 1 and gt.min() >= 0\n        assert pred.max() <= 1 and pred.min() >= 0\n        loss = F.binary_cross_entropy(pred, gt, reduction='none')\n        positive_loss = loss * positive.float()\n        negative_loss = loss * negative.float()\n\n        negative_loss, _ = torch.topk(negative_loss.view(-1), negative_count)\n\n        balance_loss = (positive_loss.sum() + negative_loss.sum()) / (\n            positive_count + negative_count + self.eps)\n\n        return balance_loss\n\n    def l1_thr_loss(self, pred, gt, mask):\n        thr_loss = torch.abs((pred - gt) * mask).sum() / (\n            mask.sum() + self.eps)\n        return thr_loss\n\n    def forward(self, preds, downsample_ratio, gt_shrink, gt_shrink_mask,\n                gt_thr, gt_thr_mask):\n        \"\"\"Compute DBNet loss.\n\n        Args:\n            preds (tensor): The output tensor with size of Nx3xHxW.\n            downsample_ratio (float): The downsample ratio for the\n                ground truths.\n            gt_shrink (list[BitmapMasks]): The mask list with each element\n                being the shrinked text mask for one img.\n            gt_shrink_mask (list[BitmapMasks]): The effective mask list with\n                each element being the shrinked effective mask for one img.\n            gt_thr (list[BitmapMasks]): The mask list with each element\n                being the threshold text mask for one img.\n            gt_thr_mask (list[BitmapMasks]): The effective mask list with\n                each element being the threshold effective mask for one img.\n\n        Returns:\n            results(dict): The dict for dbnet losses with loss_prob,\n                loss_db and loss_thresh.\n        \"\"\"\n        assert isinstance(downsample_ratio, float)\n\n        assert isinstance(gt_shrink, list)\n        assert isinstance(gt_shrink_mask, list)\n        assert isinstance(gt_thr, list)\n        assert isinstance(gt_thr_mask, list)\n\n        preds = preds[0]\n\n        pred_prob = preds[:, 0, :, :]\n        pred_thr = preds[:, 1, :, :]\n        pred_db = preds[:, 2, :, :]\n        feature_sz = preds.size()\n\n        keys = ['gt_shrink', 'gt_shrink_mask', 'gt_thr', 'gt_thr_mask']\n        gt = {}\n        for k in keys:\n            gt[k] = eval(k)\n            gt[k] = [item.rescale(downsample_ratio) for item in gt[k]]\n            gt[k] = self.bitmasks2tensor(gt[k], feature_sz[2:])\n            gt[k] = [item.to(preds.device) for item in gt[k]]\n        gt['gt_shrink'][0] = (gt['gt_shrink'][0] > 0).float()\n        if self.bbce_loss:\n            loss_prob = self.balance_bce_loss(pred_prob, gt['gt_shrink'][0],\n                                              gt['gt_shrink_mask'][0])\n        else:\n            loss_prob = self.dice_loss(pred_prob, gt['gt_shrink'][0],\n                                       gt['gt_shrink_mask'][0])\n\n        loss_db = self.dice_loss(pred_db, gt['gt_shrink'][0],\n                                 gt['gt_shrink_mask'][0])\n        loss_thr = self.l1_thr_loss(pred_thr, gt['gt_thr'][0],\n                                    gt['gt_thr_mask'][0])\n\n        results = dict(\n            loss_prob=self.alpha * loss_prob,\n            loss_db=loss_db,\n            loss_thr=self.beta * loss_thr)\n\n        return results\n", "path": "mmocr/models/textdet/losses/db_loss.py"}, {"content": "import torch\nimport torch.nn as nn\nfrom mmcv.runner import BaseModule, Sequential\nfrom mmdet.models.builder import HEADS, build_loss\n\nfrom .head_mixin import HeadMixin\n\n\n@HEADS.register_module()\nclass DBHead(HeadMixin, BaseModule):\n    \"\"\"The class for DBNet head.\n\n    This was partially adapted from https://github.com/MhLiao/DB\n    \"\"\"\n\n    def __init__(self,\n                 in_channels,\n                 with_bias=False,\n                 decoding_type='db',\n                 text_repr_type='poly',\n                 downsample_ratio=1.0,\n                 loss=dict(type='DBLoss'),\n                 train_cfg=None,\n                 test_cfg=None,\n                 init_cfg=[\n                     dict(type='Kaiming', layer='Conv'),\n                     dict(\n                         type='Constant', layer='BatchNorm', val=1., bias=1e-4)\n                 ]):\n        \"\"\"Initialization.\n\n        Args:\n            in_channels (int): The number of input channels of the db head.\n            decoding_type (str): The type of decoder for dbnet.\n            text_repr_type (str): Boundary encoding type 'poly' or 'quad'.\n            downsample_ratio (float): The downsample ratio of ground truths.\n            loss (dict): The type of loss for dbnet.\n        \"\"\"\n        super().__init__(init_cfg=init_cfg)\n\n        assert isinstance(in_channels, int)\n\n        self.in_channels = in_channels\n        self.text_repr_type = text_repr_type\n        self.loss_module = build_loss(loss)\n        self.train_cfg = train_cfg\n        self.test_cfg = test_cfg\n        self.downsample_ratio = downsample_ratio\n        self.decoding_type = decoding_type\n\n        self.binarize = Sequential(\n            nn.Conv2d(\n                in_channels, in_channels // 4, 3, bias=with_bias, padding=1),\n            nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 2, 2),\n            nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels // 4, 1, 2, 2), nn.Sigmoid())\n\n        self.threshold = self._init_thr(in_channels)\n\n    def diff_binarize(self, prob_map, thr_map, k):\n        return torch.reciprocal(1.0 + torch.exp(-k * (prob_map - thr_map)))\n\n    def forward(self, inputs):\n        prob_map = self.binarize(inputs)\n        thr_map = self.threshold(inputs)\n        binary_map = self.diff_binarize(prob_map, thr_map, k=50)\n        outputs = torch.cat((prob_map, thr_map, binary_map), dim=1)\n        return (outputs, )\n\n    def _init_thr(self, inner_channels, bias=False):\n        in_channels = inner_channels\n        seq = Sequential(\n            nn.Conv2d(\n                in_channels, inner_channels // 4, 3, padding=1, bias=bias),\n            nn.BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(inner_channels // 4, inner_channels // 4, 2, 2),\n            nn.BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(inner_channels // 4, 1, 2, 2), nn.Sigmoid())\n        return seq\n", "path": "mmocr/models/textdet/dense_heads/db_head.py"}], "after_files": [{"content": "import torch\nimport torch.nn.functional as F\nfrom mmdet.models.builder import LOSSES\nfrom torch import nn\n\nfrom mmocr.models.common.losses.dice_loss import DiceLoss\n\n\n@LOSSES.register_module()\nclass DBLoss(nn.Module):\n    \"\"\"The class for implementing DBNet loss.\n\n    This is partially adapted from https://github.com/MhLiao/DB.\n    \"\"\"\n\n    def __init__(self,\n                 alpha=1,\n                 beta=1,\n                 reduction='mean',\n                 negative_ratio=3.0,\n                 eps=1e-6,\n                 bbce_loss=False):\n        \"\"\"Initialization.\n\n        Args:\n            alpha (float): The binary loss coef.\n            beta (float): The threshold loss coef.\n            reduction (str): The way to reduce the loss.\n            negative_ratio (float): The ratio of positives to negatives.\n            eps (float): Epsilon in the threshold loss function.\n            bbce_loss (bool): Whether to use balanced bce for probability loss.\n                If False, dice loss will be used instead.\n        \"\"\"\n        super().__init__()\n        assert reduction in ['mean',\n                             'sum'], \" reduction must in ['mean','sum']\"\n        self.alpha = alpha\n        self.beta = beta\n        self.reduction = reduction\n        self.negative_ratio = negative_ratio\n        self.eps = eps\n        self.bbce_loss = bbce_loss\n        self.dice_loss = DiceLoss(eps=eps)\n\n    def bitmasks2tensor(self, bitmasks, target_sz):\n        \"\"\"Convert Bitmasks to tensor.\n\n        Args:\n            bitmasks (list[BitMasks]): The BitMasks list. Each item is for\n                one img.\n            target_sz (tuple(int, int)): The target tensor size of KxHxW\n                with K being the number of kernels.\n\n        Returns\n            result_tensors (list[tensor]): The list of kernel tensors. Each\n                element is for one kernel level.\n        \"\"\"\n        assert isinstance(bitmasks, list)\n        assert isinstance(target_sz, tuple)\n\n        batch_size = len(bitmasks)\n        num_levels = len(bitmasks[0])\n\n        result_tensors = []\n\n        for level_inx in range(num_levels):\n            kernel = []\n            for batch_inx in range(batch_size):\n                mask = torch.from_numpy(bitmasks[batch_inx].masks[level_inx])\n                mask_sz = mask.shape\n                pad = [\n                    0, target_sz[1] - mask_sz[1], 0, target_sz[0] - mask_sz[0]\n                ]\n                mask = F.pad(mask, pad, mode='constant', value=0)\n                kernel.append(mask)\n            kernel = torch.stack(kernel)\n            result_tensors.append(kernel)\n\n        return result_tensors\n\n    def balance_bce_loss(self, pred, gt, mask):\n\n        positive = (gt * mask)\n        negative = ((1 - gt) * mask)\n        positive_count = int(positive.float().sum())\n        negative_count = min(\n            int(negative.float().sum()),\n            int(positive_count * self.negative_ratio))\n\n        assert gt.max() <= 1 and gt.min() >= 0\n        assert pred.max() <= 1 and pred.min() >= 0\n        loss = F.binary_cross_entropy(pred, gt, reduction='none')\n        positive_loss = loss * positive.float()\n        negative_loss = loss * negative.float()\n\n        negative_loss, _ = torch.topk(negative_loss.view(-1), negative_count)\n\n        balance_loss = (positive_loss.sum() + negative_loss.sum()) / (\n            positive_count + negative_count + self.eps)\n\n        return balance_loss\n\n    def l1_thr_loss(self, pred, gt, mask):\n        thr_loss = torch.abs((pred - gt) * mask).sum() / (\n            mask.sum() + self.eps)\n        return thr_loss\n\n    def forward(self, preds, downsample_ratio, gt_shrink, gt_shrink_mask,\n                gt_thr, gt_thr_mask):\n        \"\"\"Compute DBNet loss.\n\n        Args:\n            preds (tensor): The output tensor with size of Nx3xHxW.\n            downsample_ratio (float): The downsample ratio for the\n                ground truths.\n            gt_shrink (list[BitmapMasks]): The mask list with each element\n                being the shrinked text mask for one img.\n            gt_shrink_mask (list[BitmapMasks]): The effective mask list with\n                each element being the shrinked effective mask for one img.\n            gt_thr (list[BitmapMasks]): The mask list with each element\n                being the threshold text mask for one img.\n            gt_thr_mask (list[BitmapMasks]): The effective mask list with\n                each element being the threshold effective mask for one img.\n\n        Returns:\n            results(dict): The dict for dbnet losses with loss_prob,\n                loss_db and loss_thresh.\n        \"\"\"\n        assert isinstance(downsample_ratio, float)\n\n        assert isinstance(gt_shrink, list)\n        assert isinstance(gt_shrink_mask, list)\n        assert isinstance(gt_thr, list)\n        assert isinstance(gt_thr_mask, list)\n\n        pred_prob = preds[:, 0, :, :]\n        pred_thr = preds[:, 1, :, :]\n        pred_db = preds[:, 2, :, :]\n        feature_sz = preds.size()\n\n        keys = ['gt_shrink', 'gt_shrink_mask', 'gt_thr', 'gt_thr_mask']\n        gt = {}\n        for k in keys:\n            gt[k] = eval(k)\n            gt[k] = [item.rescale(downsample_ratio) for item in gt[k]]\n            gt[k] = self.bitmasks2tensor(gt[k], feature_sz[2:])\n            gt[k] = [item.to(preds.device) for item in gt[k]]\n        gt['gt_shrink'][0] = (gt['gt_shrink'][0] > 0).float()\n        if self.bbce_loss:\n            loss_prob = self.balance_bce_loss(pred_prob, gt['gt_shrink'][0],\n                                              gt['gt_shrink_mask'][0])\n        else:\n            loss_prob = self.dice_loss(pred_prob, gt['gt_shrink'][0],\n                                       gt['gt_shrink_mask'][0])\n\n        loss_db = self.dice_loss(pred_db, gt['gt_shrink'][0],\n                                 gt['gt_shrink_mask'][0])\n        loss_thr = self.l1_thr_loss(pred_thr, gt['gt_thr'][0],\n                                    gt['gt_thr_mask'][0])\n\n        results = dict(\n            loss_prob=self.alpha * loss_prob,\n            loss_db=loss_db,\n            loss_thr=self.beta * loss_thr)\n\n        return results\n", "path": "mmocr/models/textdet/losses/db_loss.py"}, {"content": "import torch\nimport torch.nn as nn\nfrom mmcv.runner import BaseModule, Sequential\nfrom mmdet.models.builder import HEADS, build_loss\n\nfrom .head_mixin import HeadMixin\n\n\n@HEADS.register_module()\nclass DBHead(HeadMixin, BaseModule):\n    \"\"\"The class for DBNet head.\n\n    This was partially adapted from https://github.com/MhLiao/DB\n    \"\"\"\n\n    def __init__(self,\n                 in_channels,\n                 with_bias=False,\n                 decoding_type='db',\n                 text_repr_type='poly',\n                 downsample_ratio=1.0,\n                 loss=dict(type='DBLoss'),\n                 train_cfg=None,\n                 test_cfg=None,\n                 init_cfg=[\n                     dict(type='Kaiming', layer='Conv'),\n                     dict(\n                         type='Constant', layer='BatchNorm', val=1., bias=1e-4)\n                 ]):\n        \"\"\"Initialization.\n\n        Args:\n            in_channels (int): The number of input channels of the db head.\n            decoding_type (str): The type of decoder for dbnet.\n            text_repr_type (str): Boundary encoding type 'poly' or 'quad'.\n            downsample_ratio (float): The downsample ratio of ground truths.\n            loss (dict): The type of loss for dbnet.\n        \"\"\"\n        super().__init__(init_cfg=init_cfg)\n\n        assert isinstance(in_channels, int)\n\n        self.in_channels = in_channels\n        self.text_repr_type = text_repr_type\n        self.loss_module = build_loss(loss)\n        self.train_cfg = train_cfg\n        self.test_cfg = test_cfg\n        self.downsample_ratio = downsample_ratio\n        self.decoding_type = decoding_type\n\n        self.binarize = Sequential(\n            nn.Conv2d(\n                in_channels, in_channels // 4, 3, bias=with_bias, padding=1),\n            nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 2, 2),\n            nn.BatchNorm2d(in_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels // 4, 1, 2, 2), nn.Sigmoid())\n\n        self.threshold = self._init_thr(in_channels)\n\n    def diff_binarize(self, prob_map, thr_map, k):\n        return torch.reciprocal(1.0 + torch.exp(-k * (prob_map - thr_map)))\n\n    def forward(self, inputs):\n        prob_map = self.binarize(inputs)\n        thr_map = self.threshold(inputs)\n        binary_map = self.diff_binarize(prob_map, thr_map, k=50)\n        outputs = torch.cat((prob_map, thr_map, binary_map), dim=1)\n        return outputs\n\n    def _init_thr(self, inner_channels, bias=False):\n        in_channels = inner_channels\n        seq = Sequential(\n            nn.Conv2d(\n                in_channels, inner_channels // 4, 3, padding=1, bias=bias),\n            nn.BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(inner_channels // 4, inner_channels // 4, 2, 2),\n            nn.BatchNorm2d(inner_channels // 4), nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(inner_channels // 4, 1, 2, 2), nn.Sigmoid())\n        return seq\n", "path": "mmocr/models/textdet/dense_heads/db_head.py"}]}