SwinIR源码解读

原创

已于 2022-12-06 20:44:25 修改 · 801 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#python #开发语言

于 2022-12-06 20:29:23 首次发布

Swinir源码解读

解读几个关键函数。

特征图到窗口，窗口到特征图

def window_partition(x, window_size):
    """
    特征图矩阵转窗口矩阵
    Args:
        x: (B, H, W, C)
        window_size (int): window size
    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows


def window_reverse(windows, window_size, H, W):
    """
    窗口矩阵转特征图矩阵
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image
    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // windo