WQT and DG-YOLO: towards domain generalization in underwater object detection

As Figure 1 shows, we select 8 images with different types of water quality, and use $WCT^{2}$ [15] to transfer URPC dataset to different types of water quality. The content image is from URPC’s training set and validation set. In the following section, this seven types of training set are denoted as type1 to type7 and the corresponding validation set are denoted as val_type1 to val_type7. As for type8, only the validation set is transferred to obtain val_type8 without corresponding training set. Since model will never train on type8 domain, val_type8 is to test the domain generalization capacity of model.

A review of YOLOv3. Because AUVs with a small processing unit have limited calculation capacity, the real-time detector YOLOv3 [9] is a promising choice. YOLOv3 is a one-stage object detector, using Darknet-53 as backbone. Compared with Faster R-CNN [16], YOLOv3 does not use region proposal network. It directly regresses the coordinates of bounding box and class information with a fully convolutional network. YOLOv3 divides an image into $S\times S$ cells, and each cell is responsible for the objects lie in the cell. The training losses of YOLOv3 consists of the loss of classification $L_{cls}$, the loss of coordination $L_{coord}$, loss of object $L_{obj}$ and loss of no-object $L_{noobj}$:

$$L_{yolo}=L_{cls}+\lambda_{coord}\cdot L_{coord}+L_{obj}+\lambda_{noobj}\cdot L_{noobj},$$

(1)

where $\lambda_{coord}$ and $\lambda_{noobj}$ are trade-off parameters.

Domain Invariant Module (DIM). Since DA and DG have some similarities, we modify the domain classifier proposed by [17] to apply in our DG task. Given a batch of input images as $\{X_{1},X_{2},...,X_{N}\}$ from $K$ different source domains, its corresponding domain labels are $\{d_{1},d_{2},...,d_{N}\}$, in which $N$ is the number of batch, $d_{i}\in\mathbb{R}^{K\times 1}$. Denoting that $G$ is a feature extractor and $D$ is a domain classifier, the domain loss $L_{d}$ is defined as follows:

$$L_{d}=\sum_{i}^{N}l_{CE}(D(G(X_{i})),d_{i}),$$

(2)

where $l_{CE}$ means categorical cross entropy. In application, domain label comes from WQT, and $K$ is 7 corresponding to 7 types of water quality that WQT synthesizes. Domain loss for data from original dataset is not calculated.

IRM Penalty. Inspired by recent study [18], Invariant Risk Minimization (IRM) help learn an invariant predictor across multiple domain. Given a set of training environments (same meaning as domains) $e\in\mathcal{E}_{tr}$, our final goal is to achieve good performance across a large set of unseen but related environments $\mathcal{E}_{all}$ ($\mathcal{E}_{tr}\in\mathcal{E}_{all}$). However, directly using Empirical Risk Minimization (ERM) [19] will lead to overfitting on training environment and learn spurious correlation. In order to generalize well on unseen environments, IRM is a better choice to obtain invariance:

$${min}_{\Phi:\mathcal{X}\sim\mathcal{Y}}\ \sum_{e\in\mathcal{E}_{tr}}R^{e}(\Phi)+\lambda\cdot||\nabla_{r|r=1.0}R^{e}(r\ \cdot\Phi)||^{2},$$

(3)

where $\Phi$ is the entire invariant predictors, $R^{e}(\Phi)$ is ERM term on environment $e$, $r=1.0$ is a fixed scalar, $||\nabla_{r|r=1.0}R^{e}(r\ \cdot\Phi)||^{2}$ is invariance penalty, and $\lambda\in[0,\infty)$ is a trade-off parameter balancing the ERM term and the invariance penalty. To apply IRM in YOLOv3, the IRM penalty specially for YOLOv3 is designed as follows:

		$\displaystyle Pen_{coord}=\nabla_{r|r=1.0}||\sum_{i=0}^{S^{2}}\sum_{j=0}^{B}\mathbb{I}^{obj}_{i,j}[(x_{i}-\hat{x}_{i}\cdot r)^{2}+(y_{i}-\hat{y}_{i}\cdot r)^{2},$
		$\displaystyle\ \ \ \ \ \ \ \ \ \ +(w_{i}-\hat{w}_{i}\cdot r)^{2}+(h_{i}-\hat{h}_{i}\cdot r)^{2}]||^{2},$

$$Pen_{cls}=\nabla_{r|r=1.0}||\sum_{i=0}^{S^{2}}\mathbb{I}^{obj}_{i}\sum_{c\in classes}l_{CE}(p_{i}(c),\sigma(a_{i}(c)\cdot r))||^{2},$$

$$Pen_{obj}=\nabla_{r|r=1.0}||\sum_{i=0}^{S^{2}}\sum_{j=0}^{B}\mathbb{I}^{obj}_{i,j}l_{CE}(C_{i},\sigma(S_{i}\cdot r))||^{2},$$

$$Pen_{noobj}=\nabla_{r|r=1.0}||\sum_{i=0}^{S^{2}}\sum_{j=0}^{B}\mathbb{I}^{noobj}_{i,j}l_{CE}(C_{i},\sigma(S_{i}\cdot r))||^{2},$$

$$P_{IRM}=pen_{coord}+pen_{cls}+pen_{obj}+pen_{noobj},$$

(4)

where $\mathbb{I}_{i}^{obj}$ denotes if object appears in cell $i$, $\mathbb{I}_{i,j}^{obj}$ denotes that the $j$th bounding box predictor in cell $i$ is responsible for that prediction, $\sigma$ is sigmoid operation, $\hat{a_{i}}(c)$ is the score of class $c$ before sigmoid operation, $p_{i}\in\mathbb{R}^{K\times 1}$ is class label, $\hat{S_{i}}$ is the score of objects before sigmoid operation, $C_{i}\in\{0,1\}$ is object label. $\{\hat{x_{i}},\hat{y_{i}},\hat{w_{i}},\hat{h_{i}}\}$ is the bounding box outputted by YOLOv3, whose corresponding ground truth is $\{x_{i},y_{i},w_{i},h_{i}\}$.

Penalty term is designed based on corresponding losses of YOLOv3. To be specific, $r$ is added to different places of losses. Square gradient of losses to $r$ is the corresponding penalty term.

Network overview. An overview of our network is shown in Figure 2, we denote it DG-YOLO. Compared to YOLOv3, DIM and IRM penalty are added. In details, the backbone of YOLO darknet-53 can be regarded as a feature extractor. The feature maps extracted from darknet will be fed into Gradient Reversal Layer (GRL) [17] first, which reverses the gradient when backpropagating for adversarial learning. After that, domain classifier distinguish feature maps between domains. With the help of GRL and domain classifier, the backbone will be forced to abandon information of water quality to fool domain classifier. As a result, DG-YOLO can make a prediction depending more on semantic information. Moreover, IRM penalty is calculated simultaneously with YOLO loss. Combining (1), (2) and (4) ,the total loss of DG-YOLO is:

$$L_{total}=L_{yolo}+\lambda_{d}\cdot L_{d}+\lambda_{p}\cdot P_{IRM},$$

(5)

$\lambda_{p}$ and $\lambda_{d}$ set to 1 in experiment. In inference stage, because DIM and IRM penalty can be abandoned, DG-YOLO doesn’t affect the speed of YOLOv3. It should be emphasized that because domain label comes from WQT, DG-YOLO can not be used alone without WQT.

We evaluate WQT and DG-YOLO on a publicly available datasets: URPC2019¹¹1www.cnurpc.org, which consists of 3765 training samples and 942 validation samples over five categories: echinus, starfish, holothurian, scallop and waterweeds. Applying WQT on training set and validation set of URPC2019, we can synthesize type1-7 for training and val_type1-8 for validation. The performance on val_type8 represents domain generalization capacity of model.

YOLOv3 and DG-YOLO is trained for 300 epochs and evaluated on original and all synthetic validation sets, with image resizing to 416 $\times$ 416. Models are trained on a Nvidia GTX 1080Ti GPU with PyTorch implementation, setting batch size to 8. Adam algorithm is adopted for optimization and learning rate sets to 0.001, with $\beta_{1}=0.9$ and $\beta_{2}=0.999$. IoU, confidence and non-max suppression threshold all set to 0.5. Accumulating gradient is leveraged, which is summing up the gradient and make one step of gradient descent in each two iterations. We do not use any other data augmentation on YOLOv3 and DG-YOLO unless we mention it.

In this subsection, we analyze why WQT works. In Table 1, Ori means original URPC dataset, baseline means YOLOv3 is trained only on original dataset, and ori+type1 means YOLOv3 is trained with original dataset and type1 dataset. Full_WQT means YOLOv3 is trained across type1 to type7. From Table 1, we can find three interesting points:

(1) Compared to baseline, it can be concluded that every group of augmentation improves the performance in original validation dataset. WQT can be used together with other data augmentation methods to obtain higher performance (last two rows of Table 1), which further proves its effectiveness. Besides, there is a phenomenon that WQT also helps model generalize better on other type of water quality in most of the cases. For example, ori+type7 evaluates on type3 get mAP 39.23%, 12.66% higher than baseline.

(2) We believe that there is a correlation between performance and similarity between water qualities. First, we use style loss proposed by [20] to represent style distance, and calculate the style distance between different types water quality. We feed style image type1 to type7 into $WCT^{2}$, and extract the feature maps at certain layers from both encoder and decoder, calculating style loss between any two types and obtaining $H_{dist}$. The result is shown in Table 2. Second, we take the data from column 3 to 9 (val_+type1 to val_type7) and row 2 to 8 (model ori+type1 to ori+type7) in Table 1, subtracting each row of this 7 $\times$ 7 matrix to the performance of corresponding type of baseline, getting $H_{perf}$. Using Pearson Correlation Coefficient and taking negative, it can be found that the correlation coefficient between style and performance is 0.4634. From this analysis, it can be inferred that the increase of generalization capacity gaining from WQT is from the similarity between different types of water quality.

(3) To further analyze the finding of (2), model is evaluated on val_type8 which is a very different style from type1 to type7. There is no doubt that the WQT-trained model will perform not only better on original dataset, but also across type1 to type7 dataset. However, the model still fails on val_type8 (see Table 3), which is far from the requirement of a GUOD. WQT is not enough for domain generalization.

The effectiveness of DG-YOLO. WQT helps YOLOv3 to learn domain-invariant information, but the model still suffers from domain shift severely. In Table 3, it is shown that DG-YOLO further digs domain-invariant information from data, obtaining 3.21% mAP improvement on val_type8 compared to WQT-only. Besides, compared with other object detectors on val_type8 performance, DG-YOLO shows its much better domain generalization capacity.

Ablation study. The result of ablation study is shown in Table 3. Compared to WQT-only on val_type8, WQT+DIM has 4.52% performance decrease and WQT+$P_{IRM}$ has little improvement. However, WQT+DG-YOLO achieves 3.21% improvement, which suggests only by combining DIM and $P_{IRM}$ can lead to better performance.

Visualization of DG-YOLO. One thing that can not be ignored is that there is performance decrease in original validation dataset of WQT+DG-YOLO. It is because WQT-only is ”cheating”, learning spurious correlation to make predictions. For example, waterweeds are green in greenish water, but they may turn black in another type of water. Therefore, color of the object is not a domain-invariant information, although it is convenient to use this spurious correlation to achieve good result in just one domain. The performance decrease of DG-YOLO can be interpreted that the model abandons the domain-related information and tries to learn domain invariant information from dataset. We use SmoothGrad [21] visualization technique to prove our hypothesis, finding the area that make model to believe there is echinus with probability higher than 95%. As is shown in Figure 3, baseline focuses on the shadow on the top left of image where there is no echinus. The pixel that WQT-only focuses is too dispersed, which means WQT-only learns spurious correlation. And the pixel DG-YOLO focuses is concentrated and exactly lie on the place where there is echinus. The visualization shows that DG-YOLO learn more semantic information than baseline and WQT-only.