[Improve] Add exception for PointRend for support CPU-only usage (#1271)

* [Improve] Add exception for PointRend for support CPU-only usage

* fixed linting
pull/1801/head
Jerry Jiarui XU 2022-02-09 00:52:56 -05:00 committed by GitHub
parent 548eec9580
commit b405ce0506
1 changed files with 8 additions and 1 deletions

View File

@ -4,7 +4,11 @@
import torch
import torch.nn as nn
from mmcv.cnn import ConvModule
from mmcv.ops import point_sample
try:
from mmcv.ops import point_sample
except ModuleNotFoundError:
point_sample = None
from mmseg.models.builder import HEADS
from mmseg.ops import resize
@ -75,6 +79,9 @@ class PointHead(BaseCascadeDecodeHead):
init_cfg=dict(
type='Normal', std=0.01, override=dict(name='fc_seg')),
**kwargs)
if point_sample is None:
raise RuntimeError('Please install mmcv-full for '
'point_sample ops')
self.num_fcs = num_fcs
self.coarse_pred_each_layer = coarse_pred_each_layer