[NNCF] Enable fp16 optimization

mmseg/apis/test.py

@@ -63,7 +63,6 @@ def single_gpu_test(model,
     Returns:
         list: The prediction results.
     """
-
     model.eval()
 
     results = []

mmseg/apis/train.py

@@ -18,7 +18,7 @@
 
 import numpy as np
 import torch
-from mmcv.runner import HOOKS, build_optimizer, build_runner
+from mmcv.runner import HOOKS, build_optimizer, build_runner, Fp16OptimizerHook
 from mmcv.utils import build_from_cfg
 
 from mmseg.core import (
@@ -86,7 +86,7 @@ def needed_collect_dataset_stat(cfg):
             losses = head.loss_decode
             if not isinstance(losses, (tuple, list)):
                 losses = [losses]
-        
+
             for loss in losses:
                 if loss.type == 'MarginCalibrationLoss':
                     return True
@@ -234,7 +234,12 @@ def train_segmentor(model,
     if 'type' not in cfg.optimizer_config:
         optimizer_config = CustomOptimizerHook(**cfg.optimizer_config)
     else:
-        optimizer_config = cfg.optimizer_config
+        if cfg.optimizer_config['type'] == 'Fp16OptimizerHook':
+            config = cfg.optimizer_config
+            config.pop('type')
+            optimizer_config = Fp16OptimizerHook(**config)
+        else:
+            optimizer_config = cfg.optimizer_config
 
     # register EMA hook
     ema_cfg = cfg.get('ema_config', None)

mmseg/models/decode_heads/decode_head.py

@@ -90,7 +90,7 @@ def __init__(self,
         self.act_cfg = act_cfg
         self.ignore_index = ignore_index
         self.align_corners = align_corners
-        self.fp16_enabled = False
+        self.fp16_enabled = True
         self.enable_out_norm = enable_out_norm
         self.enable_loss_equalizer = enable_loss_equalizer
         self.loss_target = loss_target

mmseg/models/necks/fpn.py

@@ -82,7 +82,7 @@ def __init__(self,
         self.num_outs = num_outs
         self.relu_before_extra_convs = relu_before_extra_convs
         self.no_norm_on_lateral = no_norm_on_lateral
-        self.fp16_enabled = False
+        self.fp16_enabled = True
         self.upsample_cfg = upsample_cfg.copy()
 
         if end_level == -1:

mmseg/models/segmentors/base.py

@@ -26,7 +26,7 @@ class BaseSegmentor(nn.Module):
     def __init__(self):
         super(BaseSegmentor, self).__init__()
 
-        self.fp16_enabled = False
+        self.fp16_enabled = True
 
     @property
     def with_neck(self):