import numpy as np

import megengine as mge
import megengine.functional as F
import megengine.module as M
import megengine.optimizer as optim

class AdaptiveAvgPool2d(M.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return F.mean(F.mean(x, axis=-2, keepdims=True), axis=-1, keepdims=True)

class AdaptiveMaxPool2d(M.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return F.max(F.max(x, axis=-2, keepdims=True), axis=-1, keepdims=True)

class ChannelAttention(M.Module):
    def __init__(self, in_planes, ratio=16):
        super().__init__()
        self.avg_pool = AdaptiveAvgPool2d()
        self.max_pool = AdaptiveMaxPool2d()
        self.sharedMLP = M.Sequential(
            M.Conv2d(in_planes, in_planes // ratio, 1, bias=False), M.ReLU(),
            
M.Conv2d(in_planes // ratio, in_planes, 1, bias=False)