pytorch-image-models

1
""" PyTorch selectable adaptive pooling
2
Adaptive pooling with the ability to select the type of pooling from:
3
    * 'avg' - Average pooling
4
    * 'max' - Max pooling
5
    * 'avgmax' - Sum of average and max pooling re-scaled by 0.5
6
    * 'avgmaxc' - Concatenation of average and max pooling along feature dim, doubles feature dim
7

8
Both a functional and a nn.Module version of the pooling is provided.
9

10
Hacked together by / Copyright 2020 Ross Wightman
11
"""
12
from typing import Optional, Tuple, Union
13

14
import torch
15
import torch.nn as nn
16
import torch.nn.functional as F
17

18
from .format import get_spatial_dim, get_channel_dim
19

20
_int_tuple_2_t = Union[int, Tuple[int, int]]
21

22

23
def adaptive_pool_feat_mult(pool_type='avg'):
24
    if pool_type.endswith('catavgmax'):
25
        return 2
26
    else:
27
        return 1
28

29

30
def adaptive_avgmax_pool2d(x, output_size: _int_tuple_2_t = 1):
31
    x_avg = F.adaptive_avg_pool2d(x, output_size)
32
    x_max = F.adaptive_max_pool2d(x, output_size)
33
    return 0.5 * (x_avg + x_max)
34

35

36
def adaptive_catavgmax_pool2d(x, output_size: _int_tuple_2_t = 1):
37
    x_avg = F.adaptive_avg_pool2d(x, output_size)
38
    x_max = F.adaptive_max_pool2d(x, output_size)
39
    return torch.cat((x_avg, x_max), 1)
40

41

42
def select_adaptive_pool2d(x, pool_type='avg', output_size: _int_tuple_2_t = 1):
43
    """Selectable global pooling function with dynamic input kernel size
44
    """
45
    if pool_type == 'avg':
46
        x = F.adaptive_avg_pool2d(x, output_size)
47
    elif pool_type == 'avgmax':
48
        x = adaptive_avgmax_pool2d(x, output_size)
49
    elif pool_type == 'catavgmax':
50
        x = adaptive_catavgmax_pool2d(x, output_size)
51
    elif pool_type == 'max':
52
        x = F.adaptive_max_pool2d(x, output_size)
53
    else:
54
        assert False, 'Invalid pool type: %s' % pool_type
55
    return x
56

57

58
class FastAdaptiveAvgPool(nn.Module):
59
    def __init__(self, flatten: bool = False, input_fmt: F = 'NCHW'):
60
        super(FastAdaptiveAvgPool, self).__init__()
61
        self.flatten = flatten
62
        self.dim = get_spatial_dim(input_fmt)
63

64
    def forward(self, x):
65
        return x.mean(self.dim, keepdim=not self.flatten)
66

67

68
class FastAdaptiveMaxPool(nn.Module):
69
    def __init__(self, flatten: bool = False, input_fmt: str = 'NCHW'):
70
        super(FastAdaptiveMaxPool, self).__init__()
71
        self.flatten = flatten
72
        self.dim = get_spatial_dim(input_fmt)
73

74
    def forward(self, x):
75
        return x.amax(self.dim, keepdim=not self.flatten)
76

77

78
class FastAdaptiveAvgMaxPool(nn.Module):
79
    def __init__(self, flatten: bool = False, input_fmt: str = 'NCHW'):
80
        super(FastAdaptiveAvgMaxPool, self).__init__()
81
        self.flatten = flatten
82
        self.dim = get_spatial_dim(input_fmt)
83

84
    def forward(self, x):
85
        x_avg = x.mean(self.dim, keepdim=not self.flatten)
86
        x_max = x.amax(self.dim, keepdim=not self.flatten)
87
        return 0.5 * x_avg + 0.5 * x_max
88

89

90
class FastAdaptiveCatAvgMaxPool(nn.Module):
91
    def __init__(self, flatten: bool = False, input_fmt: str = 'NCHW'):
92
        super(FastAdaptiveCatAvgMaxPool, self).__init__()
93
        self.flatten = flatten
94
        self.dim_reduce = get_spatial_dim(input_fmt)
95
        if flatten:
96
            self.dim_cat = 1
97
        else:
98
            self.dim_cat = get_channel_dim(input_fmt)
99

100
    def forward(self, x):
101
        x_avg = x.mean(self.dim_reduce, keepdim=not self.flatten)
102
        x_max = x.amax(self.dim_reduce, keepdim=not self.flatten)
103
        return torch.cat((x_avg, x_max), self.dim_cat)
104

105

106
class AdaptiveAvgMaxPool2d(nn.Module):
107
    def __init__(self, output_size: _int_tuple_2_t = 1):
108
        super(AdaptiveAvgMaxPool2d, self).__init__()
109
        self.output_size = output_size
110

111
    def forward(self, x):
112
        return adaptive_avgmax_pool2d(x, self.output_size)
113

114

115
class AdaptiveCatAvgMaxPool2d(nn.Module):
116
    def __init__(self, output_size: _int_tuple_2_t = 1):
117
        super(AdaptiveCatAvgMaxPool2d, self).__init__()
118
        self.output_size = output_size
119

120
    def forward(self, x):
121
        return adaptive_catavgmax_pool2d(x, self.output_size)
122

123

124
class SelectAdaptivePool2d(nn.Module):
125
    """Selectable global pooling layer with dynamic input kernel size
126
    """
127
    def __init__(
128
            self,
129
            output_size: _int_tuple_2_t = 1,
130
            pool_type: str = 'fast',
131
            flatten: bool = False,
132
            input_fmt: str = 'NCHW',
133
    ):
134
        super(SelectAdaptivePool2d, self).__init__()
135
        assert input_fmt in ('NCHW', 'NHWC')
136
        self.pool_type = pool_type or ''  # convert other falsy values to empty string for consistent TS typing
137
        if not pool_type:
138
            self.pool = nn.Identity()  # pass through
139
            self.flatten = nn.Flatten(1) if flatten else nn.Identity()
140
        elif pool_type.startswith('fast') or input_fmt != 'NCHW':
141
            assert output_size == 1, 'Fast pooling and non NCHW input formats require output_size == 1.'
142
            if pool_type.endswith('catavgmax'):
143
                self.pool = FastAdaptiveCatAvgMaxPool(flatten, input_fmt=input_fmt)
144
            elif pool_type.endswith('avgmax'):
145
                self.pool = FastAdaptiveAvgMaxPool(flatten, input_fmt=input_fmt)
146
            elif pool_type.endswith('max'):
147
                self.pool = FastAdaptiveMaxPool(flatten, input_fmt=input_fmt)
148
            else:
149
                self.pool = FastAdaptiveAvgPool(flatten, input_fmt=input_fmt)
150
            self.flatten = nn.Identity()
151
        else:
152
            assert input_fmt == 'NCHW'
153
            if pool_type == 'avgmax':
154
                self.pool = AdaptiveAvgMaxPool2d(output_size)
155
            elif pool_type == 'catavgmax':
156
                self.pool = AdaptiveCatAvgMaxPool2d(output_size)
157
            elif pool_type == 'max':
158
                self.pool = nn.AdaptiveMaxPool2d(output_size)
159
            else:
160
                self.pool = nn.AdaptiveAvgPool2d(output_size)
161
            self.flatten = nn.Flatten(1) if flatten else nn.Identity()
162

163
    def is_identity(self):
164
        return not self.pool_type
165

166
    def forward(self, x):
167
        x = self.pool(x)
168
        x = self.flatten(x)
169
        return x
170

171
    def feat_mult(self):
172
        return adaptive_pool_feat_mult(self.pool_type)
173

174
    def __repr__(self):
175
        return self.__class__.__name__ + '(' \
176
               + 'pool_type=' + self.pool_type \
177
               + ', flatten=' + str(self.flatten) + ')'
178

179