Source code for gluoncv.model_zoo.fcn

"""Fully Convolutional Network with Strdie of 8"""
from __future__ import division
from mxnet.gluon import nn
import mxnet.ndarray as F
from mxnet.context import cpu
from mxnet.gluon.nn import HybridBlock
from .segbase import SegBaseModel
# pylint: disable=unused-argument,abstract-method,missing-docstring

__all__ = ['FCN', 'get_fcn', 'get_fcn_voc_resnet50', 'get_fcn_voc_resnet101']

[docs]class FCN(SegBaseModel): r"""Fully Convolutional Networks for Semantic Segmentation Parameters ---------- nclass : int Number of categories for the training dataset. backbone : string Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', 'resnet101' or 'resnet152'). norm_layer : object Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`; Reference: Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks for semantic segmentation." *CVPR*, 2015 Examples -------- >>> model = FCN(nclass=21, backbone='resnet50') >>> print(model) """ # pylint: disable=arguments-differ def __init__(self, nclass, backbone='resnet50', norm_layer=nn.BatchNorm, aux=True, **kwargs): super(FCN, self).__init__(nclass, aux, backbone, norm_layer=norm_layer, **kwargs) with self.name_scope(): self.head = _FCNHead(2048, nclass, norm_layer=norm_layer, **kwargs) self.head.initialize() self.head.collect_params().setattr('lr_mult', 10) if self.aux: self.auxlayer = _FCNHead(1024, nclass, norm_layer=norm_layer, **kwargs) self.auxlayer.initialize() self.auxlayer.collect_params().setattr('lr_mult', 10)
[docs] def forward(self, x): _, _, H, W = x.shape c3, c4 = self.base_forward(x) outputs = [] x = self.head(c4) x = F.contrib.BilinearResize2D(x, height=H, width=W) outputs.append(x) if self.aux: auxout = self.auxlayer(c3) auxout = F.contrib.BilinearResize2D(auxout, height=H, width=W) outputs.append(auxout) return tuple(outputs) else: return x
class _FCNHead(HybridBlock): # pylint: disable=redefined-outer-name def __init__(self, in_channels, channels, norm_layer, **kwargs): super(_FCNHead, self).__init__() with self.name_scope(): self.block = nn.HybridSequential() inter_channels = in_channels // 4 with self.block.name_scope(): self.block.add(nn.Conv2D(in_channels=in_channels, channels=inter_channels, kernel_size=3, padding=1)) self.block.add(norm_layer(in_channels=inter_channels)) self.block.add(nn.Activation('relu')) self.block.add(nn.Dropout(0.1)) self.block.add(nn.Conv2D(in_channels=inter_channels, channels=channels, kernel_size=1)) # pylint: disable=arguments-differ def hybrid_forward(self, F, x): return self.block(x)
[docs]def get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=False, root='~/.mxnet/models', ctx=cpu(0), **kwargs): r"""FCN model from the paper `"Fully Convolutional Network for semantic segmentation" <https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_ Parameters ---------- dataset : str, default pascal_voc The dataset that model pretrained on. (pascal_voc, ade20k) pretrained : bool, default False Whether to load the pretrained weights for model. ctx : Context, default CPU The context in which to load the pretrained weights. root : str, default '~/.mxnet/models' Location for keeping the model parameters. Examples -------- >>> model = get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=False) >>> print(model) """ acronyms = { 'pascal_voc': 'voc', 'ade20k': 'ade', } # infer number of classes from ..data.segbase import get_segmentation_dataset data = get_segmentation_dataset(dataset) model = FCN(data.num_class, backbone=backbone, **kwargs) if pretrained: from .model_store import get_model_file model.load_params(get_model_file('fcn_%s_%s'%(backbone, acronyms[dataset]), root=root), ctx=ctx) return model
[docs]def get_fcn_voc_resnet50(**kwargs): r"""FCN model with base network ResNet-50 pre-trained on Pascal VOC dataset from the paper `"Fully Convolutional Network for semantic segmentation" <https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_ Parameters ---------- pretrained : bool, default False Whether to load the pretrained weights for model. ctx : Context, default CPU The context in which to load the pretrained weights. root : str, default '~/.mxnet/models' Location for keeping the model parameters. Examples -------- >>> model = get_fcn_voc_resnet50(pretrained=True) >>> print(model) """ return get_fcn('pascal_voc', 'resnet50', **kwargs)
[docs]def get_fcn_voc_resnet101(**kwargs): r"""FCN model with base network ResNet-101 pre-trained on Pascal VOC dataset from the paper `"Fully Convolutional Network for semantic segmentation" <https://people.eecs.berkeley.edu/~jonlong/long_shelhamer_fcn.pdf>`_ Parameters ---------- pretrained : bool, default False Whether to load the pretrained weights for model. ctx : Context, default CPU The context in which to load the pretrained weights. root : str, default '~/.mxnet/models' Location for keeping the model parameters. Examples -------- >>> model = get_fcn_voc_resnet101(pretrained=True) >>> print(model) """ return get_fcn('pascal_voc', 'resnet101', **kwargs)