def __init__(self, inp, oup, stride, expand_ratio):
    super(InvertedResidual, self).__init__()
    self.stride = stride
    assert (stride in [1, 2])
    hidden_dim = int(round((inp * expand_ratio)))
    self.use_res_connect = ((self.stride == 1) and (inp == oup))
    self.conv = []
    if (expand_ratio == 1):
        self.conv.append(nn.MaxPool2d(kernel_size=(3, 3), stride=stride, padding=1))
        self.conv.append(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False))
        self.conv.append(nn.BatchNorm2d(oup))
    else:
        self.conv.append(nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False))
        self.conv.append(nn.BatchNorm2d(hidden_dim))
        self.conv.append(nn.PReLU())
        self.conv.append(nn.MaxPool2d(kernel_size=(3, 3), stride=stride, padding=1))
        self.conv.append(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False))
        self.conv.append(nn.BatchNorm2d(oup))
    self.conv = nn.Sequential(*self.conv)
def __init__(self, inp, oup, stride, expand_ratio):
    super(InvertedResidual, self).__init__()
    self.stride = stride
    assert (stride in [1, 2])
    hidden_dim = int(round((inp * expand_ratio)))
    self.use_res_connect = ((self.stride == 1) and (inp == oup))
    self.conv = []
    if (expand_ratio == 1):
        self.conv.append(nn.MaxPool2d(kernel_size=(3, 3), stride=stride, padding=1))
        self.conv.append(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False))
        self.conv.append(nn.BatchNorm2d(oup))
    else:
        self.conv.append(nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False))
        self.conv.append(nn.BatchNorm2d(hidden_dim))
        self.conv.append(nn.PReLU())
        self.conv.append(nn.MaxPool2d(kernel_size=(3, 3), stride=stride, padding=1))
        self.conv.append(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False))
        self.conv.append(nn.BatchNorm2d(oup))
    self.conv = nn.Sequential(*self.conv)