models/Attention_SSN.py

from abc import abstractmethod
from functools import partial
from typing import Iterable
import math

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

from .SSN import  Conv, Conv2DMod, Decoder, Up
from .attention import AttentionBlock
from .blocks import ResBlock, Res_Type, get_activation


class Attention_Encoder(nn.Module):
    def __init__(self, in_channels=3, mid_act='gelu', dropout=0.0, num_heads=8, resnet=True):
        super(Attention_Encoder, self).__init__()

        self.in_conv        = Conv(in_channels, 32-in_channels, stride=1, activation=mid_act, resnet=resnet)
        self.down_32_64     = Conv(32, 64, stride=2, activation=mid_act, resnet=resnet)
        self.down_64_64_1   = Conv(64, 64, activation=mid_act, resnet=resnet)

        self.down_64_128    = Conv(64, 128, stride=2, activation=mid_act, resnet=resnet)
        self.down_128_128_1 = Conv(128, 128,  activation=mid_act, resnet=resnet)

        self.down_128_256   = Conv(128, 256, stride=2, activation=mid_act, resnet=resnet)
        self.down_256_256_1 = Conv(256, 256, activation=mid_act, resnet=resnet)
        self.down_256_256_1_attn = AttentionBlock(256, num_heads)

        self.down_256_512   = Conv(256, 512, stride=2, activation=mid_act, resnet=resnet)
        self.down_512_512_1 = Conv(512, 512, activation=mid_act, resnet=resnet)
        self.down_512_512_1_attn = AttentionBlock(512, num_heads)

        self.down_512_512_2 = Conv(512, 512, activation=mid_act, resnet=resnet)
        self.down_512_512_2_attn = AttentionBlock(512, num_heads)

        self.down_512_512_3 = Conv(512, 512, activation=mid_act, resnet=resnet)
        self.down_512_512_3_attn = AttentionBlock(512, num_heads)


    def forward(self, x):
        x1 = self.in_conv(x)  # 32 x 256 x 256
        x1 = torch.cat((x, x1), dim=1)

        x2 = self.down_32_64(x1)
        x3 = self.down_64_64_1(x2)

        x4 = self.down_64_128(x3)
        x5 = self.down_128_128_1(x4)

        x6 = self.down_128_256(x5)
        x7 = self.down_256_256_1(x6)
        x7 = self.down_256_256_1_attn(x7)

        x8 = self.down_256_512(x7)
        x9 = self.down_512_512_1(x8)
        x9 = self.down_512_512_1_attn(x9)

        x10 = self.down_512_512_2(x9)
        x10 = self.down_512_512_2_attn(x10)

        x11 = self.down_512_512_3(x10)
        x11 = self.down_512_512_3_attn(x11)

        return x11, x10, x9, x8, x7, x6, x5, x4, x3, x2, x1


class Attention_Decoder(nn.Module):
    def __init__(self, out_channels=3, mid_act='gelu', out_act='sigmoid', resnet = True, num_heads=8):

        super(Attention_Decoder, self).__init__()

        input_channel = 512
        fea_dim       = 100

        self.to_style1 = nn.Linear(in_features=fea_dim, out_features=input_channel)

        self.up_16_16_1 = Conv(input_channel, 256, activation=mid_act, style=True, resnet=resnet)
        self.up_16_16_1_attn = AttentionBlock(256, num_heads=num_heads)

        self.up_16_16_2 = Conv(768, 512, activation=mid_act, resnet=resnet)
        self.up_16_16_2_attn = AttentionBlock(512, num_heads=num_heads)

        self.up_16_16_3      = Conv(1024, 512, activation=mid_act, resnet=resnet)
        self.up_16_16_3_attn = AttentionBlock(512, num_heads=num_heads)

        self.up_16_32        = Up(1024, 256, activation=mid_act, resnet=resnet)
        self.to_style2       = nn.Linear(in_features=fea_dim, out_features=512)
        self.up_32_32_1      = Conv(512, 256, activation=mid_act, style=True, resnet=resnet)
        self.up_32_32_1_attn = AttentionBlock(256, num_heads=num_heads)

        self.up_32_64   = Up(512, 128, activation=mid_act, resnet=resnet)
        self.to_style3  = nn.Linear(in_features=fea_dim, out_features=256)
        self.up_64_64_1 = Conv(256, 128, activation=mid_act, style=True, resnet=resnet)

        self.up_64_128    = Up(256, 64, activation=mid_act, resnet=resnet)
        self.to_style4    = nn.Linear(in_features=fea_dim, out_features=128)
        self.up_128_128_1 = Conv(128, 64, activation=mid_act, style=True, resnet=resnet)

        self.up_128_256 = Up(128, 32, activation=mid_act, resnet=resnet)
        self.out_conv   = Conv(64, out_channels, activation=out_act)
        self.out_act = get_activation(out_act)


    def forward(self, x, style):
        x11, x10, x9, x8, x7, x6, x5, x4, x3, x2, x1 = x

        style1 = self.to_style1(style)
        y = self.up_16_16_1(x11, style1)  # 256 x 16 x 16
        y = self.up_16_16_1_attn(y)

        y = torch.cat((x10, y), dim=1)  # 768 x 16 x 16
        y = self.up_16_16_2(y, y)  # 512 x 16 x 16
        y = self.up_16_16_2_attn(y)


        y = torch.cat((x9, y), dim=1)  # 1024 x 16 x 16
        y = self.up_16_16_3(y, y)  # 512 x 16 x 16
        y = self.up_16_16_3_attn(y)

        y = torch.cat((x8, y), dim=1)  # 1024 x 16 x 16
        y = self.up_16_32(y, y)  # 256 x 32 x 32

        y = torch.cat((x7, y), dim=1)
        style2 = self.to_style2(style)
        y = self.up_32_32_1(y, style2)  # 256 x 32 x 32
        y = self.up_32_32_1_attn(y)

        y = torch.cat((x6, y), dim=1)
        y = self.up_32_64(y, y)

        y = torch.cat((x5, y), dim=1)
        style3 = self.to_style3(style)

        y = self.up_64_64_1(y, style3)  # 128 x 64 x 64

        y = torch.cat((x4, y), dim=1)
        y = self.up_64_128(y, y)

        y = torch.cat((x3, y), dim=1)
        style4 = self.to_style4(style)
        y = self.up_128_128_1(y, style4)  # 64 x 128 x 128

        y = torch.cat((x2, y), dim=1)
        y = self.up_128_256(y, y)  # 32 x 256 x 256

        y = torch.cat((x1, y), dim=1)
        y = self.out_conv(y, y)  # 3 x 256 x 256
        y = self.out_act(y)
        return y



class Attention_SSN(nn.Module):
    def __init__(self, in_channels, out_channels, num_heads=8, resnet=True, mid_act='gelu', out_act='gelu'):
        super(Attention_SSN, self).__init__()
        self.encoder = Attention_Encoder(in_channels, mid_act, num_heads, resnet)
        self.decoder = Attention_Decoder(out_channels, mid_act, out_act, resnet)


    def forward(self, x, softness):
        latent  = self.encoder(x)
        pred    = self.decoder(latent, softness)

        return pred


def get_model_size(model):
    param_size = 0
    import pdb; pdb.set_trace()
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()

    buffer_size = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()

    size_all_mb = (param_size + buffer_size) / 1024 ** 2
    print('model size: {:.3f}MB'.format(size_all_mb))
    # return param_size + buffer_size
    return size_all_mb


if __name__ == '__main__':
    model = AttentionBlock(in_channels=256, num_heads=8)
    x = torch.randn(5, 256, 64, 64)

    y = model(x)
    print('{}, {}'.format(x.shape, y.shape))

    # ------------------------------------------------------------------ #
    in_channels  = 3
    out_channels = 1
    num_heads    = 8
    resnet       = True
    mid_act      = 'gelu'
    out_act      = 'gelu'

    model = Attention_SSN(in_channels=in_channels,
                           out_channels=out_channels,
                           num_heads=num_heads,
                           resnet=resnet,
                           mid_act=mid_act,
                           out_act=out_act)

    x        = torch.randn(5, 3, 256, 256)
    softness = torch.randn(5, 100)


    y = model(x, softness)


    print('x: {}, y: {}'.format(x.shape, y.shape))

    get_model_size(model)
    # ------------------------------------------------------------------ #