File size: 4,094 Bytes
320c5b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | # LeGR-Conv FiLM2 with BASNet Refinement Module
## Overview
The `legrconv4x_film2_refine.py` model extends the FiLM2 architecture by integrating the refinement module (RefUnet) from BASNet (Boundary-Aware Salient Object Detection). This enhancement improves boundary accuracy for both segmentation masks and grasp predictions using an efficient two-module approach.
## Key Components
### 1. **RefUnet (from BASNet)**
The refinement module directly adapted from BASNet:
- Takes coarse predictions as input (1 or 4 channels)
- Processes through encoder-decoder structure with skip connections
- Learns residual corrections
- Returns refined prediction by adding residual to input: `output = input + residual`
### 2. **Multi-Task Adaptation**
Our model uses an efficient two-module refinement approach:
- `refine_mask`: RefUnet for instance segmentation mask refinement (1-channel input/output)
- `refine_grasp`: RefUnet for joint grasp prediction refinement
- Takes concatenated 4-channel input: [quality, sin, cos, width]
- Learns correlations between grasp components
- Outputs refined 4-channel predictions that are split back
- Reduces model parameters from 5 to 2 RefUnet modules
### 3. **Training Strategy**
Following BASNet's multi-output supervision approach:
- **Refined predictions**: Primary loss with full weight (1.0)
- **Coarse predictions**: Auxiliary loss with weight 0.5
- This helps the network learn better representations at both stages
## Architecture Details
### RefUnet Structure (from BASNet)
```
Input (Coarse Prediction) [B, C, H, W] # C=1 for mask, C=4 for grasp
↓
Conv0 → Conv1 → Pool1
↓ ↓(skip)
Conv2 → Pool2
↓ ↓(skip)
Conv3 → Pool3
↓ ↓(skip)
Conv4 → Pool4
↓ ↓(skip)
Conv5
↓
Upsample → Concat with Conv4 skip → Conv_d4
↓
Upsample → Concat with Conv3 skip → Conv_d3
↓
Upsample → Concat with Conv2 skip → Conv_d2
↓
Upsample → Concat with Conv1 skip → Conv_d1
↓
Conv_d0 (Residual) [B, C, H, W]
↓
Output = Input + Residual
```
### Integration with FiLM2
1. **Coarse Predictions**: Generated through the standard FiLM2 pipeline
2. **Refinement**:
- Mask refined independently through `refine_mask`
- Grasp predictions concatenated and jointly refined through `refine_grasp`
3. **Loss Computation**: Both coarse and refined predictions are supervised
## Usage
### Training
```bash
python train_network_refactored.py \
--config config/OCID-VLG/legrconv_film2_refine.yaml
```
### Configuration
Key parameters in the config:
```yaml
TRAIN:
network_name: legrconv4x_film2_refine
use_refinement: True # Enable/disable refinement modules
```
### Disabling Refinement (for ablation)
```yaml
TRAIN:
use_refinement: False # Will use coarse predictions only
```
## Implementation Highlights
### Key Differences from Original BASNet
1. **Two RefUnets**: One for mask, one for joint grasp refinement (reduced from 5 modules)
2. **Input Channels**:
- Mask RefUnet: 1-channel input/output
- Grasp RefUnet: 4-channel input/output (concatenated predictions)
3. **Auxiliary Loss Weight**: 0.5 for coarse predictions (adjustable)
### Loss Computation
```python
# Refined prediction losses (primary)
loss_refined = compute_loss(refined_pred, target)
# Coarse prediction losses (auxiliary)
loss_coarse = 0.5 * compute_loss(coarse_pred, target)
# Total loss
total_loss = loss_refined + loss_coarse
```
## Expected Improvements
Based on BASNet's performance:
- **Better boundary accuracy**: RefUnet focuses on boundary refinement
- **Smoother predictions**: Residual learning helps preserve details
- **More stable training**: Multi-level supervision improves convergence
- **Efficient refinement**: Joint grasp refinement leverages correlations between components
- **Reduced parameters**: Only 2 RefUnet modules instead of 5, making training more efficient
## References
- BASNet: Qin et al., "BASNet: Boundary-Aware Salient Object Detection", CVPR 2019
- Original implementation: https://github.com/xuebinqin/BASNet |