customize.py 4.49 KB
Newer Older
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
1
2
3
4
5
6
7
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## 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
Hang Zhang's avatar
sync BN  
Hang Zhang committed
8
## LICENSE file in the root directory of this source tree
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
9
10
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

Hang Zhang's avatar
sync BN  
Hang Zhang committed
11
"""Encoding Custermized NN Module"""
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
12
import torch
Hang Zhang's avatar
sync BN  
Hang Zhang committed
13
from torch.nn import Module, Sequential, Conv2d, ReLU, AdaptiveAvgPool2d
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
14
15
16
17
from torch.nn import functional as F

from .syncbn import BatchNorm2d

Hang Zhang's avatar
path  
Hang Zhang committed
18
__all__ = ['GramMatrix', 'View', 'Sum', 'Mean', 'Normalize', 'PyramidPooling']
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
19
20


Hang Zhang's avatar
path  
Hang Zhang committed
21
22
class GramMatrix(Module):
    r""" Gram Matrix for a 4D convolutional featuremaps as a mini-batch
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
23
24

    .. math::
Hang Zhang's avatar
path  
Hang Zhang committed
25
        \mathcal{G} = \sum_{h=1}^{H_i}\sum_{w=1}^{W_i} \mathcal{F}_{h,w}\mathcal{F}_{h,w}^T
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
26
    """
Hang Zhang's avatar
path  
Hang Zhang committed
27
28
29
30
31
32
    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
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46


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):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
47
        return input.view(self.size)
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
48
49


Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
50
51
52
53
54
55
56
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):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
57
        return input.sum(self.dim, self.keep_dim)
Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
58
59
60
61
62
63
64
65
66


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):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
67
        return input.mean(self.dim, self.keep_dim)
Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
68
69


Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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
Hang Zhang's avatar
sync BN  
Hang Zhang committed
92
        self.dim = dim
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
93
94

    def forward(self, x):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
95
        return F.normalize(x, self.p, self.dim, eps=1e-10)
Hang Zhang's avatar
v1.0.1  
Hang Zhang committed
96
97


Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
98
99
class PyramidPooling(Module):
    """
Hang Zhang's avatar
sync BN  
Hang Zhang committed
100
    Reference:
Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
    """
    def __init__(self, in_channels):
        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),
                                BatchNorm2d(out_channels),
                                ReLU(True))
        self.conv2 = Sequential(Conv2d(in_channels, out_channels, 1),
                                BatchNorm2d(out_channels),
                                ReLU(True))
        self.conv3 = Sequential(Conv2d(in_channels, out_channels, 1),
                                BatchNorm2d(out_channels),
                                ReLU(True))
        self.conv4 = Sequential(Conv2d(in_channels, out_channels, 1),
                                BatchNorm2d(out_channels),
                                ReLU(True))

    def _cat_each(self, x, feat1, feat2, feat3, feat4):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
125
        assert(len(x) == len(feat1))
Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
126
127
        z = []
        for i in range(len(x)):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
128
            z.append(torch.cat((x[i], feat1[i], feat2[i], feat3[i], feat4[i]), 1))
Hang Zhang's avatar
v0.1.0  
Hang Zhang committed
129
130
131
        return z

    def forward(self, x):
Hang Zhang's avatar
sync BN  
Hang Zhang committed
132
133
134
135
136
137
        _, _, h, w = x.size()
        feat1 = F.upsample(self.conv1(self.pool1(x)), (h, w), mode='bilinear')
        feat2 = F.upsample(self.conv2(self.pool2(x)), (h, w), mode='bilinear')
        feat3 = F.upsample(self.conv3(self.pool3(x)), (h, w), mode='bilinear')
        feat4 = F.upsample(self.conv4(self.pool4(x)), (h, w), mode='bilinear')
        return torch.cat((x, feat1, feat2, feat3, feat4), 1)