CoAtNet

举报
AI浩 发表于 2021/12/22 23:15:14 2021/12/22
【摘要】 import torch import torch.nn as nn from einops import rearrange from einops.layers.torch import Rearr... 
import torch
import torch.nn as nn

from einops import rearrange
from einops.layers.torch import Rearrange


def conv_3x3_bn(inp, oup, image_size, downsample=False):
    stride = 1 if downsample == False else 2
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        nn.GELU()
    )


class PreNorm(nn.Module):
    def __init__(self, dim, fn, norm):
        super().__init__()
        self.norm = norm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class SE(nn.Module):
    def __init__(self, inp, oup, expansion=0.25):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(oup, int(inp * expansion), bias=False),
            nn.GELU(),
            nn.Linear(int(inp * expansion), oup, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)


class MBConv(nn.Module):
    def __init__(self, inp, oup, image_size, downsample=False, expansion=4):
        super().__init__()
        self.downsample = downsample
        stride = 1 if self.downsample == False else 2
        hidden_dim = int(inp * expansion)

        if self.downsample:
            self.pool = nn.MaxPool2d(3, 2, 1)
            self.proj = nn.Conv2d(inp, oup, 1, 1, 0, bias=False)

        if expansion == 1:
            self.conv = nn.Sequential(
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride,
                          1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.GELU(),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:
            self.conv = nn.Sequential(
                # pw
                # down-sample in the first conv
                nn.Conv2d(inp, hidden_dim, 1, stride, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.GELU(),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, 1, 1,
                          groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.GELU(),
                SE(inp, hidden_dim),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        
        self.conv = PreNorm(inp, self.conv, nn.BatchNorm2d)

    def forward(self, x):
        if self.downsample:
            return self.proj(self.pool(x)) + self.conv(x)
        else:
            return x + self.conv(x)


class Attention(nn.Module):
    def __init__(self, inp, oup, image_size, heads=8, dim_head=32, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not (heads == 1 and dim_head == inp)

        self.ih, self.iw = image_size

        self.heads = heads
        self.scale = dim_head ** -0.5

        # parameter table of relative position bias
        self.relative_bias_table = nn.Parameter(
            torch.zeros((2 * self.ih - 1) * (2 * self.iw - 1), heads))

        coords = torch.meshgrid((torch.arange(self.ih), torch.arange(self.iw)))
        coords = torch.flatten(torch.stack(coords), 1)
        relative_coords = coords[:, :, None] - coords[:, None, :]

        relative_coords[0] += self.ih - 1
        relative_coords[1] += self.iw - 1
        relative_coords[0] *= 2 * self.iw - 1
        relative_coords = rearrange(relative_coords, 'c h w -> h w c')
        relative_index = relative_coords.sum(-1).flatten().unsqueeze(1)
        self.register_buffer("relative_index", relative_index)

        self.attend = nn.Softmax(dim=-1)
        self.to_qkv = nn.Linear(inp, inner_dim * 3, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, oup),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(
            t, 'b n (h d) -> b h n d', h=self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale

        # Use "gather" for more efficiency on GPUs
        relative_bias = self.relative_bias_table.gather(
            0, self.relative_index.repeat(1, self.heads))
        relative_bias = rearrange(
            relative_bias, '(h w) c -> 1 c h w', h=self.ih*self.iw, w=self.ih*self.iw)
        dots = dots + relative_bias

        attn = self.attend(dots)
        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out = self.to_out(out)
        return out


class Transformer(nn.Module):
    def __init__(self, inp, oup, image_size, heads=8, dim_head=32, downsample=False, dropout=0.):
        super().__init__()
        hidden_dim = int(inp * 4)

        self.ih, self.iw = image_size
        self.downsample = downsample

        if self.downsample:
            self.pool1 = nn.MaxPool2d(3, 2, 1)
            self.pool2 = nn.MaxPool2d(3, 2, 1)
            self.proj = nn.Conv2d(inp, oup, 1, 1, 0, bias=False)

        self.attn = Attention(inp, oup, image_size, heads, dim_head, dropout)
        self.ff = FeedForward(oup, hidden_dim, dropout)

        self.attn = nn.Sequential(
            Rearrange('b c ih iw -> b (ih iw) c'),
            PreNorm(inp, self.attn, nn.LayerNorm),
            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)
        )

        self.ff = nn.Sequential(
            Rearrange('b c ih iw -> b (ih iw) c'),
            PreNorm(oup, self.ff, nn.LayerNorm),
            Rearrange('b (ih iw) c -> b c ih iw', ih=self.ih, iw=self.iw)
        )

    def forward(self, x):
        if self.downsample:
            x = self.proj(self.pool1(x)) + self.attn(self.pool2(x))
        else:
            x = x + self.attn(x)
        x = x + self.ff(x)
        return x


class CoAtNet(nn.Module):
    def __init__(self, image_size, in_channels, num_blocks, channels, num_classes=1000, block_types=['C', 'C', 'T', 'T']):
        super().__init__()
        ih, iw = image_size
        block = {'C': MBConv, 'T': Transformer}

        self.s0 = self._make_layer(
            conv_3x3_bn, in_channels, channels[0], num_blocks[0], (ih // 2, iw // 2))
        self.s1 = self._make_layer(
            block[block_types[0]], channels[0], channels[1], num_blocks[1], (ih // 4, iw // 4))
        self.s2 = self._make_layer(
            block[block_types[1]], channels[1], channels[2], num_blocks[2], (ih // 8, iw // 8))
        self.s3 = self._make_layer(
            block[block_types[2]], channels[2], channels[3], num_blocks[3], (ih // 16, iw // 16))
        self.s4 = self._make_layer(
            block[block_types[3]], channels[3], channels[4], num_blocks[4], (ih // 32, iw // 32))

        self.pool = nn.AvgPool2d(ih // 32, 1)
        self.fc = nn.Linear(channels[-1], num_classes, bias=False)

    def forward(self, x):
        x = self.s0(x)
        x = self.s1(x)
        x = self.s2(x)
        x = self.s3(x)
        x = self.s4(x)

        x = self.pool(x).view(-1, x.shape[1])
        x = self.fc(x)
        return x

    def _make_layer(self, block, inp, oup, depth, image_size):
        layers = nn.ModuleList([])
        for i in range(depth):
            if i == 0:
                layers.append(block(inp, oup, image_size, downsample=True))
            else:
                layers.append(block(oup, oup, image_size))
        return nn.Sequential(*layers)


def coatnet_0():
    num_blocks = [2, 2, 3, 5, 2]            # L
    channels = [64, 96, 192, 384, 768]      # D
    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)


def coatnet_1():
    num_blocks = [2, 2, 6, 14, 2]           # L
    channels = [64, 96, 192, 384, 768]      # D
    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)


def coatnet_2():
    num_blocks = [2, 2, 6, 14, 2]           # L
    channels = [128, 128, 256, 512, 1026]   # D
    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)


def coatnet_3():
    num_blocks = [2, 2, 6, 14, 2]           # L
    channels = [192, 192, 384, 768, 1536]   # D
    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)


def coatnet_4():
    num_blocks = [2, 2, 12, 28, 2]          # L
    channels = [192, 192, 384, 768, 1536]   # D
    return CoAtNet((224, 224), 3, num_blocks, channels, num_classes=1000)


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == '__main__':
    img = torch.randn(1, 3, 224, 224)

    net = coatnet_0()
    out = net(img)
    print(out.shape, count_parameters(net))

    net = coatnet_1()
    out = net(img)
    print(out.shape, count_parameters(net))

    net = coatnet_2()
    out = net(img)
    print(out.shape, count_parameters(net))

    net = coatnet_3()
    out = net(img)
    print(out.shape, count_parameters(net))

    net = coatnet_4()
    out = net(img)
    print(out.shape, count_parameters(net))


  
 
    
  
  • 1
    • 
  
  • 2
    • 
  
  • 3
    • 
  
  • 4
    • 
  
  • 5
    • 
  
  • 6
    • 
  
  • 7
    • 
  
  • 8
    • 
  
  • 9
    • 
  
  • 10
    • 
  
  • 11
    • 
  
  • 12
    • 
  
  • 13
    • 
  
  • 14
    • 
  
  • 15
    • 
  
  • 16
    • 
  
  • 17
    • 
  
  • 18
    • 
  
  • 19
    • 
  
  • 20
    • 
  
  • 21
    • 
  
  • 22
    • 
  
  • 23
    • 
  
  • 24
    • 
  
  • 25
    • 
  
  • 26
    • 
  
  • 27
    • 
  
  • 28
    • 
  
  • 29
    • 
  
  • 30
    • 
  
  • 31
    • 
  
  • 32
    • 
  
  • 33
    • 
  
  • 34
    • 
  
  • 35
    • 
  
  • 36
    • 
  
  • 37
    • 
  
  • 38
    • 
  
  • 39
    • 
  
  • 40
    • 
  
  • 41
    • 
  
  • 42
    • 
  
  • 43
    • 
  
  • 44
    • 
  
  • 45
    • 
  
  • 46
    • 
  
  • 47
    • 
  
  • 48
    • 
  
  • 49
    • 
  
  • 50
    • 
  
  • 51
    • 
  
  • 52
    • 
  
  • 53
    • 
  
  • 54
    • 
  
  • 55
    • 
  
  • 56
    • 
  
  • 57
    • 
  
  • 58
    • 
  
  • 59
    • 
  
  • 60
    • 
  
  • 61
    • 
  
  • 62
    • 
  
  • 63
    • 
  
  • 64
    • 
  
  • 65
    • 
  
  • 66
    • 
  
  • 67
    • 
  
  • 68
    • 
  
  • 69
    • 
  
  • 70
    • 
  
  • 71
    • 
  
  • 72
    • 
  
  • 73
    • 
  
  • 74
    • 
  
  • 75
    • 
  
  • 76
    • 
  
  • 77
    • 
  
  • 78
    • 
  
  • 79
    • 
  
  • 80
    • 
  
  • 81
    • 
  
  • 82
    • 
  
  • 83
    • 
  
  • 84
    • 
  
  • 85
    • 
  
  • 86
    • 
  
  • 87
    • 
  
  • 88
    • 
  
  • 89
    • 
  
  • 90
    • 
  
  • 91
    • 
  
  • 92
    • 
  
  • 93
    • 
  
  • 94
    • 
  
  • 95
    • 
  
  • 96
    • 
  
  • 97
    • 
  
  • 98
    • 
  
  • 99
    • 
  
  • 100
    • 
  
  • 101
    • 
  
  • 102
    • 
  
  • 103
    • 
  
  • 104
    • 
  
  • 105
    • 
  
  • 106
    • 
  
  • 107
    • 
  
  • 108
    • 
  
  • 109
    • 
  
  • 110
    • 
  
  • 111
    • 
  
  • 112
    • 
  
  • 113
    • 
  
  • 114
    • 
  
  • 115
    • 
  
  • 116
    • 
  
  • 117
    • 
  
  • 118
    • 
  
  • 119
    • 
  
  • 120
    • 
  
  • 121
    • 
  
  • 122
    • 
  
  • 123
    • 
  
  • 124
    • 
  
  • 125
    • 
  
  • 126
    • 
  
  • 127
    • 
  
  • 128
    • 
  
  • 129
    • 
  
  • 130
    • 
  
  • 131
    • 
  
  • 132
    • 
  
  • 133
    • 
  
  • 134
    • 
  
  • 135
    • 
  
  • 136
    • 
  
  • 137
    • 
  
  • 138
    • 
  
  • 139
    • 
  
  • 140
    • 
  
  • 141
    • 
  
  • 142
    • 
  
  • 143
    • 
  
  • 144
    • 
  
  • 145
    • 
  
  • 146
    • 
  
  • 147
    • 
  
  • 148
    • 
  
  • 149
    • 
  
  • 150
    • 
  
  • 151
    • 
  
  • 152
    • 
  
  • 153
    • 
  
  • 154
    • 
  
  • 155
    • 
  
  • 156
    • 
  
  • 157
    • 
  
  • 158
    • 
  
  • 159
    • 
  
  • 160
    • 
  
  • 161
    • 
  
  • 162
    • 
  
  • 163
    • 
  
  • 164
    • 
  
  • 165
    • 
  
  • 166
    • 
  
  • 167
    • 
  
  • 168
    • 
  
  • 169
    • 
  
  • 170
    • 
  
  • 171
    • 
  
  • 172
    • 
  
  • 173
    • 
  
  • 174
    • 
  
  • 175
    • 
  
  • 176
    • 
  
  • 177
    • 
  
  • 178
    • 
  
  • 179
    • 
  
  • 180
    • 
  
  • 181
    • 
  
  • 182
    • 
  
  • 183
    • 
  
  • 184
    • 
  
  • 185
    • 
  
  • 186
    • 
  
  • 187
    • 
  
  • 188
    • 
  
  • 189
    • 
  
  • 190
    • 
  
  • 191
    • 
  
  • 192
    • 
  
  • 193
    • 
  
  • 194
    • 
  
  • 195
    • 
  
  • 196
    • 
  
  • 197
    • 
  
  • 198
    • 
  
  • 199
    • 
  
  • 200
    • 
  
  • 201
    • 
  
  • 202
    • 
  
  • 203
    • 
  
  • 204
    • 
  
  • 205
    • 
  
  • 206
    • 
  
  • 207
    • 
  
  • 208
    • 
  
  • 209
    • 
  
  • 210
    • 
  
  • 211
    • 
  
  • 212
    • 
  
  • 213
    • 
  
  • 214
    • 
  
  • 215
    • 
  
  • 216
    • 
  
  • 217
    • 
  
  • 218
    • 
  
  • 219
    • 
  
  • 220
    • 
  
  • 221
    • 
  
  • 222
    • 
  
  • 223
    • 
  
  • 224
    • 
  
  • 225
    • 
  
  • 226
    • 
  
  • 227
    • 
  
  • 228
    • 
  
  • 229
    • 
  
  • 230
    • 
  
  • 231
    • 
  
  • 232
    • 
  
  • 233
    • 
  
  • 234
    • 
  
  • 235
    • 
  
  • 236
    • 
  
  • 237
    • 
  
  • 238
    • 
  
  • 239
    • 
  
  • 240
    • 
  
  • 241
    • 
  
  • 242
    • 
  
  • 243
    • 
  
  • 244
    • 
  
  • 245
    • 
  
  • 246
    • 
  
  • 247
    • 
  
  • 248
    • 
  
  • 249
    • 
  
  • 250
    • 
  
  • 251
    • 
  
  • 252
    • 
  
  • 253
    • 
  
  • 254
    • 
  
  • 255
    • 
  
  • 256
    • 
  
  • 257
    • 
  
  • 258
    • 
  
  • 259
    • 
  
  • 260
    • 
  
  • 261
    • 
  
  • 262
    • 
  
  • 263
    • 
  
  • 264
    • 
  
  • 265
    • 
  
  • 266
    • 
  
  • 267
    • 
  
  • 268
    • 
  
  • 269
    • 
  
  • 270
    • 
  
  • 271
    • 
  
  • 272
    • 
  
  • 273
    • 
  
  • 274
    • 
  
  • 275
    • 
  
  • 276
    • 
  
  • 277
    • 
  
  • 278
    • 
  
  • 279
    • 
  
  • 280
    • 
  
  • 281
    • 
  
  • 282
    • 
  
  • 283
    • 
  
  • 284
    • 
  
  • 285
    • 
  
  • 286
    • 
  
  • 287
    • 
  
  • 288
    • 
  
  • 289
    • 
  
  • 290
    • 
  
  • 291
    • 
  
  • 292
    • 
  
  • 293
    • 
  
  • 294
    • 
  
  • 295
    • 
  
  • 296
    • 
  
  • 297
    • 
  
  • 298
    •

文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。

原文链接：wanghao.blog.csdn.net/article/details/121861275

【版权声明】本文为华为云社区用户转载文章，如果您发现本社区中有涉嫌抄袭的内容，欢迎发送邮件进行举报，并提供相关证据，一经查实，本社区将立刻删除涉嫌侵权内容，举报邮箱： cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论（0）

0/1000
抱歉，系统识别当前为高风险访问，暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称，即可参与社区互动！

*长度不超过10个汉字或20个英文字符，设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符，设置后3个月内不可修改。

确认 取消

加入云驻计划，成为创作者

  • 华为云周边好礼
  • 免费体验产品
  • 特殊身份标识
  • 线下官方门票
  • 内部专家零距离
  • 与10000+优质创作者共同成长
立即加入