如果你正在处理高分辨率图像,比如医学影像分析或卫星图像识别,是否经常遇到显存不足的问题?传统的注意力机制在处理长序列时计算复杂度呈平方级增长,这让很多CV任务在现实场景中难以落地。而滑动窗口注意力通过局部注意力机制将复杂度降低到线性级别,但现有的实现往往忽略了像素级别的精细依赖关系。
最近CCF-A类期刊TIFS 2026年发表的重构滑动窗口注意力RSWAtt,正是针对这一痛点提出的创新解决方案。它不仅继承了滑动窗口注意力的高效特性,更重要的是通过重构机制实现了真正的像素级依赖建模,在保持计算效率的同时大幅提升了视觉任务的精度。
本文将从实际应用角度深入解析RSWAtt的核心原理,提供完整的代码实现,并展示如何在常见的CV任务中即插即用这一先进技术。无论你是从事目标检测、图像分割还是其他视觉任务,都能从中获得实用的技术方案。
1. RSWAtt要解决的核心问题
1.1 传统注意力机制在CV任务中的瓶颈
在计算机视觉领域,随着图像分辨率的不断提升,传统的自注意力机制面临着严峻的计算挑战。一个1080p的高清图像包含超过200万个像素点,如果使用标准的Transformer架构,注意力矩阵的大小将达到(200万×200万),这在实际应用中是完全不可行的。
滑动窗口注意力通过将全局注意力限制在局部窗口内,确实解决了计算复杂度的问题。但这种方法引入了新的问题:窗口之间的信息交互不足,难以捕捉长距离的像素依赖关系。特别是在边缘检测、细粒度分割等任务中,这种局限性尤为明显。
1.2 RSWAtt的创新突破点
RSWAtt的核心创新在于引入了"重构机制",它不仅仅是在局部窗口内计算注意力,更重要的是通过多层次的特征重构来实现窗口间的信息融合。这种设计使得模型能够在保持线性计算复杂度的同时,建立更加丰富的像素级依赖关系。
具体来说,RSWAtt通过三个关键改进实现了这一目标:
- 动态窗口重构:根据图像内容自适应调整窗口大小和形状
- 跨窗口信息传递:设计专门的机制促进不同窗口间的特征交互
- 多尺度特征融合:在不同尺度上建立像素依赖关系
2. 滑动窗口注意力的基础概念
2.1 标准自注意力机制回顾
在深入理解RSWAtt之前,我们需要回顾一下标准的自注意力机制。自注意力的计算公式为:
import torch import torch.nn as nn import torch.nn.functional as F class StandardSelfAttention(nn.Module): def __init__(self, dim, heads=8, dim_head=None): super().__init__() self.heads = heads self.dim_head = dim_head if dim_head else dim // heads self.scale = self.dim_head ** -0.5 self.to_qkv = nn.Linear(dim, self.dim_head * heads * 3, bias=False) self.to_out = nn.Linear(self.dim_head * heads, dim) def forward(self, x): b, n, d = x.shape qkv = self.to_qkv(x).chunk(3, dim=-1) q, k, v = map(lambda t: t.reshape(b, n, self.heads, self.dim_head).transpose(1, 2), qkv) dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale attn = dots.softmax(dim=-1) out = torch.matmul(attn, v) out = out.transpose(1, 2).reshape(b, n, -1) return self.to_out(out)这种标准实现的计算复杂度为O(n²),其中n是序列长度。对于高分辨率图像,这种复杂度是无法接受的。
2.2 滑动窗口注意力的基本原理
滑动窗口注意力通过限制每个位置只能关注其邻近的位置来降低计算复杂度。具体实现如下:
class SlidingWindowAttention(nn.Module): def __init__(self, dim, window_size, heads=8, dim_head=None): super().__init__() self.heads = heads self.dim_head = dim_head if dim_head else dim // heads self.scale = self.dim_head ** -0.5 self.window_size = window_size self.to_qkv = nn.Linear(dim, self.dim_head * heads * 3, bias=False) self.to_out = nn.Linear(self.dim_head * heads, dim) def create_window_masks(self, h, w, device): """创建滑动窗口掩码""" masks = [] for i in range(0, h, self.window_size): for j in range(0, w, self.window_size): mask = torch.zeros(h, w, device=device) i_end = min(i + self.window_size, h) j_end = min(j + self.window_size, w) mask[i:i_end, j:j_end] = 1 masks.append(mask) return torch.stack(masks) def forward(self, x, h, w): b, n, d = x.shape assert n == h * w, "序列长度必须等于高度×宽度" # 将序列重塑为空间格式 x_spatial = x.reshape(b, h, w, d) # 应用滑动窗口注意力 output_windows = [] window_masks = self.create_window_masks(h, w, x.device) for mask in window_masks: window_indices = mask.nonzero(as_tuple=False) if len(window_indices) == 0: continue # 提取窗口内的特征 window_features = x_spatial[:, window_indices[:, 0], window_indices[:, 1], :] window_features = window_features.reshape(b, -1, d) # 在窗口内应用标准注意力 qkv = self.to_qkv(window_features).chunk(3, dim=-1) q, k, v = map(lambda t: t.reshape(b, -1, self.heads, self.dim_head).transpose(1, 2), qkv) dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale attn = dots.softmax(dim=-1) window_out = torch.matmul(attn, v) window_out = window_out.transpose(1, 2).reshape(b, -1, self.dim_head * self.heads) window_out = self.to_out(window_out) # 将结果放回原位置 output_window = torch.zeros(b, n, d, device=x.device) flat_indices = window_indices[:, 0] * w + window_indices[:, 1] output_window[:, flat_indices] = window_out output_windows.append(output_window) # 合并所有窗口的结果 output = sum(output_windows) / len(output_windows) return output这种实现的复杂度降低到了O(n×w²),其中w是窗口大小,通常远小于序列长度n。
3. RSWAtt的核心原理与架构设计
3.1 重构机制的核心思想
RSWAtt的核心创新在于引入了特征重构机制。传统的滑动窗口注意力在独立的窗口内计算注意力,窗口之间缺乏有效的信息交互。RSWAtt通过多层次的重构操作,实现了窗口间的特征融合。
重构机制包含三个关键步骤:
- 局部特征提取:在每个窗口内进行精细的特征学习
- 跨窗口信息传递:通过重叠窗口或特征重采样实现信息交换
- 全局特征重构:整合所有窗口的信息,重建全局特征表示
3.2 RSWAtt的完整架构实现
下面是RSWAtt的完整PyTorch实现:
import torch import torch.nn as nn import torch.nn.functional as F class RSWAtt(nn.Module): def __init__(self, dim, window_size, shift_size=0, num_heads=8, qkv_bias=True, attn_drop=0., proj_drop=0.): super().__init__() self.dim = dim self.window_size = window_size self.shift_size = shift_size self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) # 重构相关的参数 self.reconstruction_mlp = nn.Sequential( nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim) ) self.softmax = nn.Softmax(dim=-1) def forward(self, x, mask=None): """ Args: x: input features with shape of (num_windows*B, N, C) mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None """ B_, N, C = x.shape qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # 形状: (B_, num_heads, N, head_dim) q = q * self.scale attn = (q @ k.transpose(-2, -1)) # (B_, num_heads, N, N) if mask is not None: nW = mask.shape[0] attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) attn = attn.view(-1, self.num_heads, N, N) attn = self.softmax(attn) attn = self.attn_drop(attn) x = (attn @ v).transpose(1, 2).reshape(B_, N, C) # 特征重构步骤 reconstructed_x = self.reconstruction_mlp(x) x = x + reconstructed_x # 残差连接 x = self.proj(x) x = self.proj_drop(x) return x class WindowPartition(nn.Module): """将特征图划分为不重叠的窗口""" def __init__(self, window_size): super().__init__() self.window_size = window_size def forward(self, x): """ Args: x: (B, H, W, C) Returns: windows: (num_windows*B, window_size, window_size, C) """ B, H, W, C = x.shape x = x.view(B, H // self.window_size, self.window_size, W // self.window_size, self.window_size, C) windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.window_size, self.window_size, C) return windows class WindowReverse(nn.Module): """将窗口反转回特征图""" def __init__(self, window_size, H, W): super().__init__() self.window_size = window_size self.H = H self.W = W def forward(self, windows): """ Args: windows: (num_windows*B, window_size, window_size, C) Returns: x: (B, H, W, C) """ B = int(windows.shape[0] / (self.H * self.W / self.window_size / self.window_size)) x = windows.view(B, self.H // self.window_size, self.W // self.window_size, self.window_size, self.window_size, -1) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, self.H, self.W, -1) return x3.3 像素级依赖建模的实现细节
RSWAtt通过特殊的注意力权重计算和特征重构机制实现像素级依赖建模:
class PixelLevelDependency(nn.Module): """像素级依赖建模模块""" def __init__(self, dim, kernel_size=3): super().__init__() self.dim = dim self.kernel_size = kernel_size self.padding = kernel_size // 2 # 局部卷积用于捕捉像素间关系 self.local_conv = nn.Conv2d(dim, dim, kernel_size, padding=self.padding, groups=dim) self.norm = nn.LayerNorm(dim) self.activation = nn.GELU() def forward(self, x, H, W): """ Args: x: (B, H*W, C) Returns: x: (B, H*W, C) 增强像素依赖的特征 """ B, N, C = x.shape assert N == H * W # 转换为空间格式 x_spatial = x.transpose(1, 2).reshape(B, C, H, W) # 应用局部卷积捕捉像素关系 local_features = self.local_conv(x_spatial) local_features = local_features.reshape(B, C, -1).transpose(1, 2) # 与原始特征融合 enhanced_x = x + local_features enhanced_x = self.norm(enhanced_x) enhanced_x = self.activation(enhanced_x) return enhanced_x class EnhancedRSWAtt(nn.Module): """增强版的RSWAtt,包含像素级依赖建模""" def __init__(self, dim, window_size, num_heads, mlp_ratio=4., drop=0., attn_drop=0.): super().__init__() self.window_size = window_size self.dim = dim # 注意力模块 self.attn = RSWAtt(dim, window_size, num_heads=num_heads, attn_drop=attn_drop, proj_drop=drop) # 像素级依赖建模 self.pixel_dependency = PixelLevelDependency(dim) # MLP层 self.mlp = nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Dropout(drop), nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(drop) ) self.norm1 = nn.LayerNorm(dim) self.norm2 = nn.LayerNorm(dim) self.norm3 = nn.LayerNorm(dim) def forward(self, x, H, W): # 第一部分:窗口注意力 shortcut = x x = self.norm1(x) # 窗口划分 x_windows = window_partition(x, self.window_size, H, W) x_windows = x_windows.view(-1, self.window_size * self.window_size, self.dim) # 应用RSWAtt attn_windows = self.attn(x_windows) attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.dim) # 窗口反转 x = window_reverse(attn_windows, self.window_size, H, W) x = shortcut + x # 第二部分:像素级依赖建模 x = x + self.pixel_dependency(self.norm2(x), H, W) # 第三部分:MLP x = x + self.mlp(self.norm3(x)) return x4. 环境准备与依赖安装
4.1 基础环境配置
在开始使用RSWAtt之前,需要确保环境满足以下要求:
# 创建conda环境 conda create -n rswatt python=3.8 conda activate rswatt # 安装PyTorch(根据CUDA版本选择) pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113 # 安装其他依赖 pip install numpy matplotlib opencv-python pillow pip install timm # 图像模型库 pip install einops # 张量操作工具4.2 项目结构规划
建议的项目结构如下:
rswatt_project/ ├── models/ │ ├── __init__.py │ ├── rswatt.py # RSWAtt核心实现 │ └── backbones.py # 骨干网络集成 ├── datasets/ │ ├── __init__.py │ └── vision_datasets.py # 数据加载器 ├── configs/ │ └── default.yaml # 配置文件 ├── utils/ │ ├── logger.py # 日志工具 │ └── metrics.py # 评估指标 ├── train.py # 训练脚本 ├── eval.py # 评估脚本 └── requirements.txt # 依赖列表4.3 依赖检查脚本
创建环境检查脚本确保所有依赖正确安装:
# check_environment.py import sys import pkg_resources def check_environment(): """检查环境依赖是否满足要求""" required_packages = { 'torch': '1.12.0', 'torchvision': '0.13.0', 'numpy': '1.21.0', 'opencv-python': '4.5.0', 'timm': '0.6.0', 'einops': '0.6.0' } missing_packages = [] version_mismatch = [] for package, required_version in required_packages.items(): try: installed_version = pkg_resources.get_distribution(package).version if pkg_resources.parse_version(installed_version) < pkg_resources.parse_version(required_version): version_mismatch.append(f"{package} (需要: {required_version}, 已安装: {installed_version})") except pkg_resources.DistributionNotFound: missing_packages.append(package) if missing_packages or version_mismatch: print("环境检查失败:") if missing_packages: print(f"缺少包: {', '.join(missing_packages)}") if version_mismatch: print(f"版本不匹配: {', '.join(version_mismatch)}") return False else: print("环境检查通过!") return True if __name__ == "__main__": check_environment()5. RSWAtt在CV任务中的即插即用实现
5.1 图像分类任务集成
将RSWAtt集成到标准的Vision Transformer中:
import torch.nn as nn from einops import rearrange class RSWAttVisionTransformer(nn.Module): """基于RSWAtt的Vision Transformer""" def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., window_size=7, drop_rate=0., attn_drop_rate=0.): super().__init__() self.num_classes = num_classes self.num_features = self.embed_dim = embed_dim self.patch_size = patch_size # 图像分块嵌入 self.patch_embed = PatchEmbed( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches # 位置编码 self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) self.pos_drop = nn.Dropout(p=drop_rate) # RSWAtt blocks self.blocks = nn.ModuleList([ EnhancedRSWAtt( dim=embed_dim, window_size=window_size, num_heads=num_heads, mlp_ratio=mlp_ratio, drop=drop_rate, attn_drop=attn_drop_rate ) for i in range(depth)]) self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() # 初始化权重 self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) elif isinstance(m, nn.Parameter): nn.init.trunc_normal_(m, std=.02) def forward_features(self, x): B = x.shape[0] x = self.patch_embed(x) # (B, L, C) x = x + self.pos_embed x = self.pos_drop(x) H, W = self.patch_embed.grid_size for blk in self.blocks: x = blk(x, H, W) x = self.norm(x) return x.mean(dim=1) # 全局平均池化 def forward(self, x): x = self.forward_features(x) x = self.head(x) return x class PatchEmbed(nn.Module): """图像分块嵌入""" def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().__init__() self.img_size = (img_size, img_size) self.patch_size = (patch_size, patch_size) self.grid_size = (img_size // patch_size, img_size // patch_size) self.num_patches = self.grid_size[0] * self.grid_size[1] self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): B, C, H, W = x.shape assert H == self.img_size[0] and W == self.img_size[1], \ f"输入图像大小({H}*{W})不匹配模型期望的{self.img_size[0]}*{self.img_size[1]}" x = self.proj(x).flatten(2).transpose(1, 2) # (B, L, C) return x5.2 目标检测任务适配
将RSWAtt集成到目标检测框架中:
class RSWAttFPN(nn.Module): """基于RSWAtt的特征金字塔网络""" def __init__(self, in_channels_list, out_channels, window_sizes=[7, 7, 7]): super().__init__() self.lateral_convs = nn.ModuleList() self.fpn_convs = nn.ModuleList() self.rswatt_blocks = nn.ModuleList() for i, in_channels in enumerate(in_channels_list): lateral_conv = nn.Conv2d(in_channels, out_channels, 1) fpn_conv = nn.Conv2d(out_channels, out_channels, 3, padding=1) rswatt_block = EnhancedRSWAtt( dim=out_channels, window_size=window_sizes[i], num_heads=out_channels // 32 ) self.lateral_convs.append(lateral_conv) self.fpn_convs.append(fpn_conv) self.rswatt_blocks.append(rswatt_block) # 初始化权重 for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_uniform_(m.weight, a=1) if m.bias is not None: nn.init.constant_(m.bias, 0) def forward(self, inputs): """ Args: inputs: 不同尺度的特征图列表 Returns: outputs: 增强后的多尺度特征 """ # lateral连接 laterals = [ lateral_conv(inputs[i]) for i, lateral_conv in enumerate(self.lateral_convs) ] # 自上而下的路径 for i in range(len(laterals) - 1, 0, -1): laterals[i - 1] += F.interpolate( laterals[i], scale_factor=2, mode="nearest" ) # FPN卷积和RSWAtt增强 outputs = [] for i, (lateral, fpn_conv, rswatt_block) in enumerate( zip(laterals, self.fpn_convs, self.rswatt_blocks) ): # 应用卷积 x = fpn_conv(lateral) # 应用RSWAtt B, C, H, W = x.shape x = x.flatten(2).transpose(1, 2) # (B, H*W, C) x = rswatt_block(x, H, W) x = x.transpose(1, 2).reshape(B, C, H, W) outputs.append(x) return outputs5.3 语义分割任务应用
在语义分割任务中集成RSWAtt:
class RSWAttUNet(nn.Module): """基于RSWAtt的UNet分割网络""" def __init__(self, in_channels=3, num_classes=21, base_dim=64, window_size=7, depths=[2, 2, 2, 2]): super().__init__() self.depths = depths # 编码器 self.encoders = nn.ModuleList() self.downsamplers = nn.ModuleList() dim = base_dim for i, depth in enumerate(depths): encoder_blocks = nn.ModuleList([ RSWAttEncoderBlock(dim, window_size) for _ in range(depth) ]) self.encoders.append(encoder_blocks) if i < len(depths) - 1: downsampler = nn.Sequential( nn.Conv2d(dim, dim * 2, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(dim * 2), nn.ReLU(inplace=True) ) self.downsamplers.append(downsampler) dim *= 2 # 解码器 self.decoders = nn.ModuleList() self.upsamplers = nn.ModuleList() for i, depth in enumerate(reversed(depths[:-1])): upsampler = nn.Sequential( nn.ConvTranspose2d(dim, dim // 2, kernel_size=2, stride=2), nn.BatchNorm2d(dim // 2), nn.ReLU(inplace=True) ) self.upsamplers.append(upsampler) dim //= 2 decoder_blocks = nn.ModuleList([ RSWAttDecoderBlock(dim, window_size) for _ in range(depth) ]) self.decoders.append(decoder_blocks) # 输出层 self.output_conv = nn.Conv2d(base_dim, num_classes, kernel_size=1) def forward(self, x): # 编码路径 skips = [] for i, (encoder_blocks, downsampler) in enumerate( zip(self.encoders, self.downsamplers + [None]) ): for block in encoder_blocks: x = block(x) skips.append(x) if downsampler is not None: x = downsampler(x) # 解码路径 for i, (upsampler, decoder_blocks, skip) in enumerate( zip(self.upsamplers, self.decoders, reversed(skips[:-1])) ): x = upsampler(x) # 跳跃连接 x = torch.cat([x, skip], dim=1) for block in decoder_blocks: x = block(x) return self.output_conv(x) class RSWAttEncoderBlock(nn.Module): """RSWAtt编码器块""" def __init__(self, dim, window_size): super().__init__() self.conv = nn.Sequential( nn.Conv2d(dim, dim, 3, padding=1), nn.BatchNorm2d(dim), nn.ReLU(inplace=True) ) self.rswatt = EnhancedRSWAtt(dim, window_size, num_heads=dim//32) def forward(self, x): # 卷积路径 conv_out = self.conv(x) # RSWAtt路径 B, C, H, W = x.shape x_flat = x.flatten(2).transpose(1, 2) attn_out = self.rswatt(x_flat, H, W) attn_out = attn_out.transpose(1, 2).reshape(B, C, H, W) return conv_out + attn_out6. 完整训练流程与代码实现
6.1 训练配置与参数设置
创建完整的训练配置类:
import yaml from dataclasses import dataclass from typing import Dict, Any @dataclass class TrainingConfig: """训练配置参数""" # 数据参数 data_path: str = "./data" batch_size: int = 32 num_workers: int = 4 image_size: int = 224 # 模型参数 model_name: str = "rswatt_vit_small" num_classes: int = 1000 patch_size: int = 16 embed_dim: int = 384 depth: int = 12 num_heads: int = 6 window_size: int = 7 mlp_ratio: float = 4.0 # 优化器参数 learning_rate: float = 1e-3 weight_decay: float = 0.05 momentum: float = 0.9 optimizer: str = "adamw" # 学习率调度 scheduler: str = "cosine" warmup_epochs: int = 5 min_lr: float = 1e-6 # 训练参数 epochs: int = 300 eval_interval: int = 1 save_interval: int = 10 resume: str = "" # 恢复训练的检查点路径 # 设备参数 device: str = "cuda" if torch.cuda.is_available() else "cpu" amp: bool = True # 自动混合精度 def to_dict(self) -> Dict[str, Any]: return {k: v for k, v in self.__dict__.items() if not k.startswith('_')} @classmethod def from_yaml(cls, yaml_path: str): with open(yaml_path, 'r') as f: config_dict = yaml.safe_load(f) return cls(**config_dict) def save_to_yaml(self, yaml_path: str): with open(yaml_path, 'w') as f: yaml.dump(self.to_dict(), f, default_flow_style=False)6.2 完整的训练脚本
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torch.cuda.amp import GradScaler, autocast import time import os from pathlib import Path class RSWAttTrainer: """RSWAtt模型训练器""" def __init__(self, config: TrainingConfig): self.config = config self.device = torch.device(config.device) self.scaler = GradScaler() if config.amp else None # 创建模型 self.model = self._create_model() self.model.to(self.device) # 创建数据加载器 self.train_loader, self.val_loader = self._create_dataloaders() # 创建优化器和调度器 self.optimizer = self._create_optimizer() self.scheduler = self._create_scheduler() # 损失函数 self.criterion = nn.CrossEntropyLoss() # 训练状态 self.epoch = 0 self.best_acc = 0.0 self.train_losses = [] self.val_accuracies = [] # 创建输出目录 self.output_dir = Path("./outputs") self.output_dir.mkdir(exist_ok=True) def _create_model(self): """创建RSWAtt模型""" if self.config.model_name == "rswatt_vit_small": return RSWAttVisionTransformer( img_size=self.config.image_size, patch_size=self.config.patch_size, num_classes=self.config.num_classes, embed_dim=self.config.embed_dim, depth=self.config.depth, num_heads=self.config.num_heads, window_size=self.config.window_size, mlp_ratio=self.config.mlp_ratio ) else: raise ValueError(f"未知模型: {self.config.model_name}") def _create_dataloaders(self): """创建数据加载器""" # 这里使用伪数据集,实际项目中替换为真实数据集 from torch.utils.data import TensorDataset # 创建伪训练数据 x_train = torch.randn(1000, 3, self.config.image_size, self.config.image_size) y_train = torch.randint(0, self.config.num_classes, (1000,)) train_dataset = TensorDataset(x_train, y_train) # 创建伪验证数据 x_val = torch.randn(200, 3, self.config.image_size, self.config.image_size) y_val = torch.randint(0, self.config.num_classes, (200,)) val_dataset = TensorDataset(x_val, y_val) train_loader = DataLoader( train_dataset, batch_size=self.config.batch_size, shuffle=True, num_workers=self.config.num_workers ) val_loader = DataLoader( val_dataset, batch_size=self.config.batch_size, shuffle=False, num_workers=self.config.num_workers ) return train_loader, val_loader def _create_optimizer(self): """创建优化器"""