Commit 3e243c1a authored by Hang Zhang's avatar Hang Zhang Committed by Facebook GitHub Bot
Browse files

Patch for Quantizing PointRend model

Summary: PointRend mask doesn't work for quantization. Add a patch to disable it.

Reviewed By: wat3rBro

Differential Revision: D27800349

fbshipit-source-id: ae0268ee78b000245ebdb2edbfc679a62c85a59a
parent a95c7983
...@@ -23,6 +23,7 @@ from mobile_cv.arch.utils.quantize_utils import ( ...@@ -23,6 +23,7 @@ from mobile_cv.arch.utils.quantize_utils import (
) )
from mobile_cv.predictor.api import FuncInfo from mobile_cv.predictor.api import FuncInfo
from torch.quantization.quantize_fx import prepare_fx, prepare_qat_fx, convert_fx from torch.quantization.quantize_fx import prepare_fx, prepare_qat_fx, convert_fx
from detectron2.projects.point_rend import PointRendMaskHead
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
...@@ -146,7 +147,6 @@ def _fx_quant_prepare(self, cfg): ...@@ -146,7 +147,6 @@ def _fx_quant_prepare(self, cfg):
self.roi_heads.box_predictor.bbox_pred, qconfig self.roi_heads.box_predictor.bbox_pred, qconfig
) )
def d2_meta_arch_prepare_for_quant(self, cfg): def d2_meta_arch_prepare_for_quant(self, cfg):
model = self model = self
torch.backends.quantized.engine = cfg.QUANTIZATION.BACKEND torch.backends.quantized.engine = cfg.QUANTIZATION.BACKEND
...@@ -155,6 +155,9 @@ def d2_meta_arch_prepare_for_quant(self, cfg): ...@@ -155,6 +155,9 @@ def d2_meta_arch_prepare_for_quant(self, cfg):
if model.training if model.training
else torch.quantization.get_default_qconfig(cfg.QUANTIZATION.BACKEND) else torch.quantization.get_default_qconfig(cfg.QUANTIZATION.BACKEND)
) )
if hasattr(model, "roi_heads") and hasattr(model.roi_heads, "mask_head") and \
isinstance(model.roi_heads.mask_head, PointRendMaskHead):
model.roi_heads.mask_head.qconfig = None
logger.info("Setup the model with qconfig:\n{}".format(model.qconfig)) logger.info("Setup the model with qconfig:\n{}".format(model.qconfig))
# Modify the model for eager mode # Modify the model for eager mode
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment