UNet 改进(25)：结合Shift-Transformer结构

1.介绍

本文将详细介绍一个结合了Shift-Transformer模块的UNet网络实现，这是一种将传统卷积神经网络与自注意力机制相结合的创新架构。

网络概述

这个网络是基于经典的UNet架构，但在瓶颈层(bottleneck)加入了Shift-Transformer模块，旨在结合CNN的局部特征提取能力和Transformer的全局建模能力。

核心组件解析

1. ShiftTransformerBlock

ShiftTransformerBlock是整个网络中最具创新性的模块，它结合了卷积操作和自注意力机制：

class ShiftTransformerBlock(nn.Module):
    def __init__(self, dim, num_heads=4, shift_size=5, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.shift_size &#