Source code for encoding.nn.customize

##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Hang Zhang
## ECE Department, Rutgers University
## Email: zhang.hang@rutgers.edu
## Copyright (c) 2017
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

"""Encoding Custermized NN Module"""
import torch
from torch.nn import Module, Sequential, Conv2d, ReLU, AdaptiveAvgPool2d, \
    NLLLoss, BCELoss, CrossEntropyLoss, AvgPool2d, MaxPool2d, Parameter
from torch.nn import functional as F
from torch.autograd import Variable

torch_ver = torch.__version__[:3]

__all__ = ['GramMatrix', 'SegmentationLosses', 'View', 'Sum', 'Mean',
           'Normalize', 'PyramidPooling']

[docs]class GramMatrix(Module): r""" Gram Matrix for a 4D convolutional featuremaps as a mini-batch .. math:: \mathcal{G} = \sum_{h=1}^{H_i}\sum_{w=1}^{W_i} \mathcal{F}_{h,w}\mathcal{F}_{h,w}^T """ def forward(self, y): (b, ch, h, w) = y.size() features = y.view(b, ch, w * h) features_t = features.transpose(1, 2) gram = features.bmm(features_t) / (ch * h * w) return gram
def softmax_crossentropy(input, target, weight, size_average, ignore_index, reduce=True): return F.nll_loss(F.log_softmax(input, 1), target, weight, size_average, ignore_index, reduce) class SegmentationLosses(CrossEntropyLoss): """2D Cross Entropy Loss with Auxilary Loss""" def __init__(self, se_loss=False, se_weight=0.2, nclass=-1, aux=False, aux_weight=0.2, weight=None, size_average=True, ignore_index=-1): super(SegmentationLosses, self).__init__(weight, size_average, ignore_index) self.se_loss = se_loss self.aux = aux self.nclass = nclass self.se_weight = se_weight self.aux_weight = aux_weight self.bceloss = BCELoss(weight, size_average) def forward(self, *inputs): if not self.se_loss and not self.aux: return super(SegmentationLosses, self).forward(*inputs) elif not self.se_loss: pred1, pred2, target = tuple(inputs) loss1 = super(SegmentationLosses, self).forward(pred1, target) loss2 = super(SegmentationLosses, self).forward(pred2, target) return loss1 + self.aux_weight * loss2 elif not self.aux: pred, se_pred, target = tuple(inputs) se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred) loss1 = super(SegmentationLosses, self).forward(pred, target) loss2 = self.bceloss(F.sigmoid(se_pred), se_target) return loss1 + self.se_weight * loss2 else: pred1, se_pred, pred2, target = tuple(inputs) se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred1) loss1 = super(SegmentationLosses, self).forward(pred1, target) loss2 = super(SegmentationLosses, self).forward(pred2, target) loss3 = self.bceloss(F.sigmoid(se_pred), se_target) return loss1 + self.aux_weight * loss2 + self.se_weight * loss3 @staticmethod def _get_batch_label_vector(target, nclass): # target is a 3D Variable BxHxW, output is 2D BxnClass batch = target.size(0) tvect = Variable(torch.zeros(batch, nclass)) for i in range(batch): hist = torch.histc(target[i].cpu().data.float(), bins=nclass, min=0, max=nclass-1) vect = hist>0 tvect[i] = vect return tvect class View(Module): """Reshape the input into different size, an inplace operator, support SelfParallel mode. """ def __init__(self, *args): super(View, self).__init__() if len(args) == 1 and isinstance(args[0], torch.Size): self.size = args[0] else: self.size = torch.Size(args) def forward(self, input): return input.view(self.size) class Sum(Module): def __init__(self, dim, keep_dim=False): super(Sum, self).__init__() self.dim = dim self.keep_dim = keep_dim def forward(self, input): return input.sum(self.dim, self.keep_dim) class Mean(Module): def __init__(self, dim, keep_dim=False): super(Mean, self).__init__() self.dim = dim self.keep_dim = keep_dim def forward(self, input): return input.mean(self.dim, self.keep_dim) class Normalize(Module): r"""Performs :math:`L_p` normalization of inputs over specified dimension. Does: .. math:: v = \frac{v}{\max(\lVert v \rVert_p, \epsilon)} for each subtensor v over dimension dim of input. Each subtensor is flattened into a vector, i.e. :math:`\lVert v \rVert_p` is not a matrix norm. With default arguments normalizes over the second dimension with Euclidean norm. Args: p (float): the exponent value in the norm formulation. Default: 2 dim (int): the dimension to reduce. Default: 1 """ def __init__(self, p=2, dim=1): super(Normalize, self).__init__() self.p = p self.dim = dim def forward(self, x): return F.normalize(x, self.p, self.dim, eps=1e-8) class PyramidPooling(Module): """ Reference: Zhao, Hengshuang, et al. *"Pyramid scene parsing network."* """ def __init__(self, in_channels, norm_layer, up_kwargs): super(PyramidPooling, self).__init__() self.pool1 = AdaptiveAvgPool2d(1) self.pool2 = AdaptiveAvgPool2d(2) self.pool3 = AdaptiveAvgPool2d(3) self.pool4 = AdaptiveAvgPool2d(6) out_channels = int(in_channels/4) self.conv1 = Sequential(Conv2d(in_channels, out_channels, 1, bias=False), norm_layer(out_channels), ReLU(True)) self.conv2 = Sequential(Conv2d(in_channels, out_channels, 1, bias=False), norm_layer(out_channels), ReLU(True)) self.conv3 = Sequential(Conv2d(in_channels, out_channels, 1, bias=False), norm_layer(out_channels), ReLU(True)) self.conv4 = Sequential(Conv2d(in_channels, out_channels, 1, bias=False), norm_layer(out_channels), ReLU(True)) # bilinear upsample options self._up_kwargs = up_kwargs def forward(self, x): _, _, h, w = x.size() feat1 = F.upsample(self.conv1(self.pool1(x)), (h, w), **self._up_kwargs) feat2 = F.upsample(self.conv2(self.pool2(x)), (h, w), **self._up_kwargs) feat3 = F.upsample(self.conv3(self.pool3(x)), (h, w), **self._up_kwargs) feat4 = F.upsample(self.conv4(self.pool4(x)), (h, w), **self._up_kwargs) return torch.cat((x, feat1, feat2, feat3, feat4), 1)