这是我的模型文件model， 加入了一个WTConv2d 层，加入之后我的train文件运行会出现错误，我需要这个分组卷积，怎么修改
Traceback (most recent call last):
File "/root/task/train.py", line 268, in <module>
raise e
File "/root/task/train.py", line 239, in <module>
train_model(
File "/root/task/train.py", line 135, in train_model
masks_pred = model(images) # 前向传播，得到预测结果
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/root/task/model666.py", line 404, in forward
x1 = self.in_conv(x)
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/container.py", line 215, in forward
input = module(input)
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/root/task/model666.py", line 191, in forward
next_x_ll = self.iwt_function(curr_x)
File "/root/task/model666.py", line 81, in inverse_wavelet_transform
out = F.conv_transpose2d(coeffs, inv_filters, stride=2, padding=0, groups=1) # 转置卷积
RuntimeError: Given transposed=1, weight of size [4, 1, 2, 2], expected input[8192, 1, 112, 112] to have 4 channels, but got 1 channels instead

帮我检查这两个文件，帮我解决报错，如果还要类似的问题，都帮我找出来并解决，给我完整代码，不要省略，并给我详细的中文注释
import torch
import torch.nn as nn
import torch.nn.functional as F
import pywt # 导入 PyWavelets 库，用于小波变换
from functools import partial

=========================

定义 WTConv2d 所需的 wavelet 函数

=========================

def create_wavelet_filter(wave_type, dtype):
"""
创建小波滤波器，用于正变换和逆变换。

text
参数：
- wave_type: 小波类型，例如 'db1'。
- dtype: 数据类型。

返回：
- filters: 正变换滤波器，形状为 [4, 1, 2, 2]。
- inv_filters: 逆变换滤波器，形状为 [4, 1, 2, 2]。
"""
wavelet = pywt.Wavelet(wave_type)
dec_hi = torch.tensor(wavelet.dec_hi[::-1], dtype=dtype)
dec_lo = torch.tensor(wavelet.dec_lo[::-1], dtype=dtype)
rec_hi = torch.tensor(wavelet.rec_hi[::-1], dtype=dtype)  # 注意逆序
rec_lo = torch.tensor(wavelet.rec_lo[::-1], dtype=dtype)  # 注意逆序

# 构建正变换滤波器
filters = torch.stack([
    dec_lo.unsqueeze(0) * dec_lo.unsqueeze(1),  # LL
    dec_lo.unsqueeze(0) * dec_hi.unsqueeze(1),  # LH
    dec_hi.unsqueeze(0) * dec_lo.unsqueeze(1),  # HL
    dec_hi.unsqueeze(0) * dec_hi.unsqueeze(1)   # HH
], dim=0)  # 形状为 [4, 2, 2]

# 构建逆变换滤波器
inv_filters = torch.stack([
    rec_lo.unsqueeze(0) * rec_lo.unsqueeze(1),  # LL
    rec_lo.unsqueeze(0) * rec_hi.unsqueeze(1),  # LH
    rec_hi.unsqueeze(0) * rec_lo.unsqueeze(1),  # HL
    rec_hi.unsqueeze(0) * rec_hi.unsqueeze(1)   # HH
], dim=0)  # 形状为 [4, 2, 2]

# 将滤波器调整为 [4, 1, 2, 2]，输入通道数为 1
filters = filters.unsqueeze(1)  # [4, 1, 2, 2]
inv_filters = inv_filters.unsqueeze(1)  # [4, 1, 2, 2]

return filters, inv_filters

def wavelet_transform(x, filters):
"""
对输入 x 进行小波变换。

text
参数：
- x: 输入张量，形状为 [batch_size, channels, height, width]。
- filters: 小波滤波器，形状为 [4, 1, 2, 2]。

返回：
- out: 变换后的张量，形状为 [batch_size, channels, 4, out_height, out_width]。
"""
batch_size, channels, height, width = x.size()
x = x.reshape(-1, 1, height, width)  # 将通道维度合并到批次维度，形状为 [batch_size * channels, 1, height, width]
out = F.conv2d(x, filters, stride=2, padding=0, groups=1)  # 卷积后形状为 [batch_size * channels, 4, out_height, out_width]
out = out.reshape(batch_size, channels, filters.size(0), out.size(2), out.size(3))  # 重新调整形状
return out

def inverse_wavelet_transform(coeffs, inv_filters):
"""
对小波系数进行逆变换。

text
参数：
- coeffs: 小波系数张量，形状为 [batch_size, channels, 4, height, width]。
- inv_filters: 逆小波滤波器，形状为 [4, 1, 2, 2]。

返回：
- out: 重构后的张量，形状为 [batch_size, channels, out_height, out_width]。
"""
batch_size, channels, num_coeffs, height, width = coeffs.size()
coeffs = coeffs.reshape(-1, 1, height, width)  # 将通道维度合并到批次维度
out = F.conv_transpose2d(coeffs, inv_filters, stride=2, padding=0, groups=1)  # 转置卷积
out = out.reshape(batch_size, channels, out.size(2), out.size(3))  # 重新调整形状
return out

=========================

定义 WTConv2d 类

=========================

class _ScaleModule(nn.Module):
"""
缩放模块，用于调整张量的尺度。
"""
def init(self, dims, init_scale=1.0, init_bias=0):
super(_ScaleModule, self).init()
self.dims = dims
self.weight = nn.Parameter(torch.ones(*dims) * init_scale)
self.bias = None

text
def forward(self, x):
    return torch.mul(self.weight, x)

class WTConv2d(nn.Module):
"""
WTConv2d 层，使用级联小波分解来增加有效感受野。

text
参数：
- in_channels: 输入通道数。
- out_channels: 输出通道数（需要与输入通道数相同）。
- kernel_size: 卷积核大小。
- stride: 步幅。
- padding: 填充大小。
- bias: 是否使用偏置。
- wt_levels: 小波分解的级数。
- wt_type: 小波类型。
"""
def __init__(self, in_channels, out_channels, kernel_size=5, stride=1, padding=0, bias=True, wt_levels=1, wt_type='db1'):
    super(WTConv2d, self).__init__()

    assert in_channels == out_channels, "WTConv2d 层要求输入和输出通道数相同"

    self.in_channels = in_channels
    self.wt_levels = wt_levels
    self.stride = stride
    self.dilation = 1
    self.padding = padding

    # 创建小波滤波器（形状为 [4, 1, 2, 2]）
    self.wt_filter, self.iwt_filter = create_wavelet_filter(wt_type, torch.float)
    self.wt_filter = nn.Parameter(self.wt_filter, requires_grad=False)
    self.iwt_filter = nn.Parameter(self.iwt_filter, requires_grad=False)

    self.wt_function = partial(wavelet_transform, filters=self.wt_filter)
    self.iwt_function = partial(inverse_wavelet_transform, inv_filters=self.iwt_filter)

    # 基础卷积层
    self.base_conv = nn.Conv2d(in_channels, in_channels, kernel_size, padding=self.padding, stride=1, dilation=1, groups=in_channels, bias=bias)
    self.base_scale = _ScaleModule([1, in_channels, 1, 1])

    # 小波卷积层
    self.wavelet_convs = nn.ModuleList([
        nn.Conv2d(in_channels * 4, in_channels * 4, kernel_size, padding=self.padding, stride=1, dilation=1, groups=in_channels * 4, bias=False)
        for _ in range(self.wt_levels)
    ])
    self.wavelet_scale = nn.ModuleList([
        _ScaleModule([1, in_channels * 4, 1, 1], init_scale=0.1)
        for _ in range(self.wt_levels)
    ])

    if self.stride > 1:
        self.stride_filter = nn.Parameter(torch.ones(in_channels, 1, 1, 1), requires_grad=False)
        self.do_stride = lambda x_in: F.conv2d(x_in, self.stride_filter, bias=None, stride=self.stride, groups=in_channels)
    else:
        self.do_stride = None

def forward(self, x):
    x_ll_in_levels = []
    x_h_in_levels = []
    shapes_in_levels = []
    curr_x_ll = x

    # 小波分解
    for i in range(self.wt_levels):
        curr_shape = curr_x_ll.shape
        shapes_in_levels.append(curr_shape)
        if (curr_shape[2] % 2 > 0) or (curr_shape[3] % 2 > 0):
            curr_pads = (0, curr_shape[3] % 2, 0, curr_shape[2] % 2)
            curr_x_ll = F.pad(curr_x_ll, curr_pads)

        curr_x = self.wt_function(curr_x_ll)  # 小波变换
        curr_x_ll = curr_x[:, :, 0, :, :]  # LL 系数
        shape_x = curr_x.shape  # [batch_size, channels, 4, height, width]

        # 将小波系数展开到通道维度
        curr_x_tag = curr_x.reshape(shape_x[0], shape_x[1] * 4, shape_x[3], shape_x[4])
        # 对展开后的张量进行卷积和尺度调整
        curr_x_tag = self.wavelet_scale[i](self.wavelet_convs[i](curr_x_tag))
        # 将张量形状恢复
        curr_x_tag = curr_x_tag.reshape(shape_x)
        x_ll_in_levels.append(curr_x_tag[:, :, 0, :, :])  # LL
        x_h_in_levels.append(curr_x_tag[:, :, 1:4, :, :])  # LH、HL、HH

    next_x_ll = 0

    # 小波逆变换
    for i in range(self.wt_levels - 1, -1, -1):
        curr_x_ll = x_ll_in_levels.pop()
        curr_x_h = x_h_in_levels.pop()
        curr_shape = shapes_in_levels.pop()
        curr_x_ll = curr_x_ll + next_x_ll
        curr_x = torch.cat([curr_x_ll.unsqueeze(2), curr_x_h], dim=2)
        next_x_ll = self.iwt_function(curr_x)
        next_x_ll = next_x_ll[:, :, :curr_shape[2], :curr_shape[3]]  # 去除填充

    x_tag = next_x_ll  # 最终的小波重构结果
    assert len(x_ll_in_levels) == 0

    # 基础卷积和尺度调整
    x = self.base_scale(self.base_conv(x))
    x = x + x_tag  # 与小波结果相加

    if self.do_stride is not None:
        x = self.do_stride(x)

    return x

=========================

以下是您的模型代码，已整合 WTConv2d 层

=========================

定义双卷积模块，使用 WTConv2d 替换部分卷积层

class DoubleConv(nn.Sequential):
def init(self, in_channels, out_channels, mid_channels=None, wt_levels=1):
if mid_channels is None:
mid_channels = out_channels
layers = []

text
    # 如果输入通道数等于中间通道数，使用 WTConv2d
    if in_channels == mid_channels:
        layers.append(WTConv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False, wt_levels=wt_levels))
    else:
        layers.append(nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False))
    layers.append(nn.BatchNorm2d(mid_channels))
    layers.append(nn.ReLU(inplace=True))

    # 如果中间通道数等于输出通道数，使用 WTConv2d
    if mid_channels == out_channels:
        layers.append(WTConv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False, wt_levels=wt_levels))
    else:
        layers.append(nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False))
    layers.append(nn.BatchNorm2d(out_channels))
    layers.append(nn.ReLU(inplace=True))

    super(DoubleConv, self).__init__(*layers)

下采样模块，使用 WTConv2d 替换第一层卷积

class Down(nn.Sequential):
def init(self, in_channels, out_channels, wt_levels=1):
super(Down, self).init(
# 使用 WTConv2d 代替步幅为2的卷积层
WTConv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=1, bias=False, wt_levels=wt_levels),
nn.BatchNorm2d(in_channels),
nn.ReLU(inplace=True),
# 双卷积模块
DoubleConv(in_channels, out_channels, wt_levels=wt_levels)
)

修改后的 Neck 模块，使用 WTConv2d

class Neck(nn.Module):
def init(self, in_channels, out_channels, reduction=16, wt_levels=1):
super(Neck, self).init()
self.conv = nn.Sequential(
WTConv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False, wt_levels=wt_levels), # 3x3 WTConv2d
nn.BatchNorm2d(out_channels), # 批量归一化
nn.ReLU(inplace=True) # ReLU 激活
)
self.channel_attention = cSE(out_channels, reduction) # 增加 cSE 通道注意力机制

text
def forward(self, x):
    x = self.conv(x)  # 进行卷积操作
    x = self.channel_attention(x)  # 应用通道注意力
    return x

上采样模块（保持不变）

class Up(nn.Module):
def init(self, in_channels, out_channels, bilinear=True):
super(Up, self).init()
if bilinear:
# 使用双线性插值进行上采样
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
# 双卷积模块
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
else:
# 使用转置卷积进行上采样
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)

text
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
    x1 = self.up(x1)
    # 处理尺寸差异
    diff_y = x2.size()[2] - x1.size()[2]
    diff_x = x2.size()[3] - x1.size()[3]
    x1 = F.pad(x1, [diff_x // 2, diff_x - diff_x // 2,
                    diff_y // 2, diff_y - diff_y // 2])
    # 拼接特征
    x = torch.cat([x2, x1], dim=1)
    x = self.conv(x)
    return x

定义输出层（保持不变）

class OutConv(nn.Sequential):
def init(self, in_channels, num_classes):
super(OutConv, self).init(
nn.Conv2d(in_channels, num_classes, kernel_size=1)
)

SimSPPF模块（保持不变）

class SimSPPF(nn.Module):
def init(self, in_channels, out_channels, pool_size=(5, 9, 13)):
super(SimSPPF, self).init()
self.conv1 = nn.Conv2d(in_channels, in_channels // 2, kernel_size=1, stride=1, padding=0)
self.bn1 = nn.BatchNorm2d(in_channels // 2)
self.relu = nn.ReLU(inplace=True)

text
    self.pool1 = nn.MaxPool2d(kernel_size=pool_size[0], stride=1, padding=pool_size[0] // 2)
    self.pool2 = nn.MaxPool2d(kernel_size=pool_size[1], stride=1, padding=pool_size[1] // 2)
    self.pool3 = nn.MaxPool2d(kernel_size=pool_size[2], stride=1, padding=pool_size[2] // 2)

    self.conv2 = nn.Conv2d(in_channels // 2 * 4, out_channels, kernel_size=1, stride=1, padding=0)
    self.bn2 = nn.BatchNorm2d(out_channels)

def forward(self, x):
    x = self.relu(self.bn1(self.conv1(x)))
    pool1 = self.pool1(x)
    pool2 = self.pool2(x)
    pool3 = self.pool3(x)
    x = torch.cat([x, pool1, pool2, pool3], dim=1)
    x = self.relu(self.bn2(self.conv2(x)))
    return x

sSE模块（保持不变）

class sSE(nn.Module):
def init(self, in_channels):
super(sSE, self).init()
self.conv = nn.Conv2d(in_channels, 1, kernel_size=1, stride=1, padding=0)
self.sigmoid = nn.Sigmoid()

text
def forward(self, x):
    x_se = self.sigmoid(self.conv(x))
    return x * x_se

cSE模块（保持不变）

class cSE(nn.Module):
def init(self, in_channels, reduction=16):
super(cSE, self).init()
self.global_pool = nn.AdaptiveAvgPool2d(1)
self.fc1 = nn.Conv2d(in_channels, in_channels // reduction, kernel_size=1)
self.fc2 = nn.Conv2d(in_channels // reduction, in_channels, kernel_size=1)
self.relu = nn.ReLU(inplace=True)
self.sigmoid = nn.Sigmoid()

text
def forward(self, x):
    x_ce = self.global_pool(x)
    x_ce = self.relu(self.fc1(x_ce))
    x_ce = self.sigmoid(self.fc2(x_ce))
    return x * x_ce

scSE模块（保持不变）

class scSE(nn.Module):
def init(self, in_channels, reduction=16):
super(scSE, self).init()
self.spatial_se = sSE(in_channels)
self.channel_se = cSE(in_channels, reduction)

text
def forward(self, x):
    x_sse = self.spatial_se(x)
    x_cse = self.channel_se(x)
    return torch.max(x_sse, x_cse)

改进版UNet模型

class self_net(nn.Module):
def init(self, n_channels: int = 3, n_classes: int = 4, bilinear: bool = True, base_c: int = 64, wt_levels=1):
super(self_net, self).init()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear

text
    # 输入卷积
    self.in_conv = DoubleConv(n_channels, base_c, wt_levels=wt_levels)

    # 下采样和SimSPPF模块
    self.down1 = Down(base_c, base_c * 2, wt_levels=wt_levels)
    self.sim_sppf1 = SimSPPF(base_c * 2, base_c * 2)

    self.down2 = Down(base_c * 2, base_c * 4, wt_levels=wt_levels)
    self.sim_sppf2 = SimSPPF(base_c * 4, base_c * 4)

    self.down3 = Down(base_c * 4, base_c * 8, wt_levels=wt_levels)
    self.sim_sppf3 = SimSPPF(base_c * 8, base_c * 8)

    factor = 2 if bilinear else 1
    self.down4 = Down(base_c * 8, base_c * 16 // factor, wt_levels=wt_levels)

    # Neck层
    self.neck = Neck(base_c * 16 // factor, base_c * 16 // factor, wt_levels=wt_levels)

    # 上采样和scSE模块
    self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
    self.scse1 = scSE(base_c * 8 // factor)

    self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear)
    self.scse2 = scSE(base_c * 4 // factor)

    self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear)
    self.scse3 = scSE(base_c * 2 // factor)

    self.up4 = Up(base_c * 2, base_c, bilinear)
    self.scse4 = scSE(base_c)

    # 输出卷积
    self.out_conv = OutConv(base_c, n_classes)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # 编码路径
    x1 = self.in_conv(x)
    x2 = self.sim_sppf1(self.down1(x1))
    x3 = self.sim_sppf2(self.down2(x2))
    x4 = self.sim_sppf3(self.down3(x3))
    x5 = self.down4(x4)

    # Neck层
    x5 = self.neck(x5)

    # 解码路径
    x = self.scse1(self.up1(x5, x4))
    x = self.scse2(self.up2(x, x3))
    x = self.scse3(self.up3(x, x2))
    x = self.scse4(self.up4(x, x1))

    # 输出
    logits = self.out_conv(x)
    return logits

这是我的train文件
import argparse
import logging
import os
import random
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
from pathlib import Path
from torch import optim
from torch.utils.data import DataLoader, random_split
from tqdm import tqdm

import wandb
from evaluate import evaluate
from model666 import self_net
from utils.data_loading import BasicDataset, CarvanaDataset
from utils.dice_score import dice_loss

设置图像、掩码和检查点的目录路径

dir_img = Path('data/imgs+1') # 图像目录
dir_mask = Path('data/masks+1') # 掩码目录
dir_checkpoint = Path('checkpoints') # 检查点目录，用于保存模型

def calculate_iou(pred, target, num_classes):
"""
计算每个类别的IoU以及平均mIoU。

text
参数：
- pred: 模型预测的输出（未经过argmax处理）
- target: 真实的掩码标签
- num_classes: 类别数量

返回：
- ious: 每个类别的IoU列表
- miou: 平均mIoU
"""
ious = []
pred = pred.argmax(dim=1)  # 对模型输出取argmax，得到每个像素的预测类别
for cls in range(num_classes):
    pred_inds = (pred == cls)  # 预测为当前类别的像素索引
    target_inds = (target == cls)  # 真实为当前类别的像素索引

    intersection = (pred_inds & target_inds).sum().float().item()  # 交集像素数
    union = (pred_inds | target_inds).sum().float().item()  # 并集像素数

    if union == 0:
        ious.append(float('nan'))  # 如果当前类别在预测和真实中都不存在，IoU记为nan
    else:
        iou = intersection / union  # 计算IoU
        ious.append(iou)

# 计算平均mIoU（忽略nan值）
valid_ious = [iou for iou in ious if not torch.isnan(torch.tensor(iou))]
miou = sum(valid_ious) / len(valid_ious) if valid_ious else float('nan')
return ious, miou

def train_model(
model,
device,
epochs: int = 200,
batch_size: int = 4,
learning_rate: float = 1e-6,
val_percent: float = 0.1,
save_checkpoint: bool = True,
img_scale: float = 0.5,
amp: bool = False,
weight_decay: float = 1e-8,
momentum: float = 0.999,
gradient_clipping: float = 1.0,
):
# 1. 创建数据集
try:
dataset = CarvanaDataset(dir_img, dir_mask, img_scale)
except (AssertionError, RuntimeError, IndexError):
dataset = BasicDataset(dir_img, dir_mask, img_scale)

text
# 2. 将数据集划分为训练集和验证集
n_val = int(len(dataset) * val_percent)  # 验证集大小
n_train = len(dataset) - n_val  # 训练集大小
train_set, val_set = random_split(dataset, [n_train, n_val], generator=torch.Generator().manual_seed(0))

# 3. 创建数据加载器
loader_args = dict(batch_size=batch_size, num_workers=2, pin_memory=True)
train_loader = DataLoader(train_set, shuffle=True, **loader_args)
val_loader = DataLoader(val_set, shuffle=False, drop_last=True, **loader_args)

# 初始化日志记录
experiment = wandb.init(project='U-Net', resume='allow', anonymous='must')
experiment.config.update(
    dict(epochs=epochs, batch_size=batch_size, learning_rate=learning_rate,
         val_percent=val_percent, save_checkpoint=save_checkpoint, img_scale=img_scale, amp=amp)
)

logging.info(f'''Starting training:
Epochs: {epochs}
Batch size: {batch_size}
Learning rate: {learning_rate}
Training size: {n_train}
Validation size: {n_val}
Checkpoints: {save_checkpoint}
Device: {device.type}
Images scaling: {img_scale}
Mixed Precision: {amp}
''')

# 4. 设置优化器、损失函数、学习率调度器和梯度缩放器
optimizer = optim.RMSprop(model.parameters(),
                          lr=learning_rate, weight_decay=weight_decay, momentum=momentum, foreach=True)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=5)  # 最大化Dice得分
grad_scaler = torch.cuda.amp.GradScaler(enabled=amp)
criterion = nn.CrossEntropyLoss()  # 使用交叉熵损失
global_step = 0  # 全局训练步数

# 5. 开始训练
best_loss = float('inf')  # 初始化最优损失为正无穷

for epoch in range(1, epochs + 1):
    model.train()  # 设置模型为训练模式
    epoch_loss = 0  # 初始化每个epoch的损失
    ious = []  # 存储每个batch的mIoU
    class_ious = [0] * model.n_classes  # 初始化每个类别的IoU之和
    class_counts = [0] * model.n_classes  # 初始化每个类别的计数

    with tqdm(total=n_train, desc=f'Epoch {epoch}/{epochs}', unit='img') as pbar:
        for batch in train_loader:
            images, true_masks = batch['image'], batch['mask']

            images = images.to(device=device, dtype=torch.float32, memory_format=torch.channels_last)
            true_masks = true_masks.to(device=device, dtype=torch.long)

            with torch.autocast(device.type if device.type != 'mps' else 'cpu', enabled=amp):
                masks_pred = model(images)  # 前向传播，得到预测结果
                loss = criterion(masks_pred, true_masks)  # 计算交叉熵损失
                loss += dice_loss(
                    F.softmax(masks_pred, dim=1).float(),
                    F.one_hot(true_masks, model.n_classes).permute(0, 3, 1, 2).float(),
                    multiclass=True
                )  # 计算Dice损失

            optimizer.zero_grad(set_to_none=True)  # 清空梯度
            grad_scaler.scale(loss).backward()  # 反向传播
            grad_scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clipping)  # 梯度裁剪
            grad_scaler.step(optimizer)  # 更新参数
            grad_scaler.update()

            pbar.update(images.shape[0])  # 更新进度条
            global_step += 1  # 更新全局步数
            epoch_loss += loss.item()  # 累加损失

            # 计算IoU
            batch_ious, batch_miou = calculate_iou(masks_pred, true_masks, num_classes=model.n_classes)
            ious.append(batch_miou)  # 记录当前batch的mIoU

            # 累加每个类别的IoU和计数
            for i, iou in enumerate(batch_ious):
                if not torch.isnan(torch.tensor(iou)):
                    class_ious[i] += iou  # 累加IoU
                    class_counts[i] += 1  # 累加计数

            # 记录到wandb
            experiment.log({
                'train loss': loss.item(),
                'step': global_step,
                'epoch': epoch,
                'mIoU': batch_miou,
            })
            pbar.set_postfix(**{'loss (batch)': loss.item(), 'mIoU (batch)': batch_miou})

    # 计算整个epoch的平均mIoU
    avg_miou = sum(ious) / len(ious)
    print(f'Epoch {epoch}, Loss/train: {epoch_loss / len(train_loader)}, mIoU: {avg_miou}')

    # 计算每个类别的平均IoU
    avg_class_ious = [class_ious[i] / class_counts[i] if class_counts[i] > 0 else float('nan') for i in range(model.n_classes)]

    # 打印每个类别的IoU
    print(f'Epoch {epoch} Class IoUs:')
    for cls_idx, iou in enumerate(avg_class_ious):
        print(f'Class {cls_idx}: IoU: {iou}')

    # 保存损失最小的模型参数
    if epoch_loss < best_loss:
        best_loss = epoch_loss
        torch.save(model.state_dict(), 'best_model.pth')
        logging.info(f'Best model saved with loss: {best_loss}')

    # 保存检查点
    if save_checkpoint:
        Path(dir_checkpoint).mkdir(parents=True, exist_ok=True)  # 确保检查点目录存在
        torch.save(model.state_dict(), str(dir_checkpoint / f'checkpoint_epoch{epoch}.pth'))  # 只保存模型参数
        logging.info(f'Checkpoint {epoch} saved!')  # 记录检查点保存信息

    # 更新学习率调度器
    scheduler.step(avg_miou)

def get_args():
parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
parser.add_argument('--epochs', '-e', metavar='E', type=int, default=400, help='Number of epochs')
parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=32, help='Batch size')
parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=1e-6,
help='Learning rate', dest='lr')
parser.add_argument('--load', '-f', type=str, default=False, help='Load model from a .pth file')
parser.add_argument('--scale', '-s', type=float, default=0.5, help='Downscaling factor of the images')
parser.add_argument('--validation', '-v', dest='val', type=float, default=10.0,
help='Percent of the data that is used as validation (0-100)')
parser.add_argument('--amp', action='store_true', default=False, help='Use mixed precision')
parser.add_argument('--bilinear', action='store_true', default=False, help='Use bilinear upsampling')
parser.add_argument('--classes', '-c', type=int, default=4, help='Number of classes')

return parser.parse_args()

if name == 'main':
args = get_args()

text
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')

# 根据需要修改n_channels和n_classes
model = self_net(n_channels=3, n_classes=4)
model = model.to(memory_format=torch.channels_last)

logging.info(f'Network:\n'
             f'\t{model.n_channels} input channels\n'
             f'\t{model.n_classes} output channels (classes)\n'
             f'\t{"Bilinear" if model.bilinear else "Transposed conv"} upscaling')

if args.load:
    state_dict = torch.load(args.load, map_location=device)
    model.load_state_dict(state_dict)
    logging.info(f'Model loaded from {args.load}')

model.to(device=device)
try:
    train_model(
        model=model,
        epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.lr,
        device=device,
        img_scale=args.scale,
        val_percent=args.val / 100,
        amp=args.amp
    )
except RuntimeError as e:
    if 'out of memory' in str(e):
        logging.error(
            'Detected OutOfMemoryError! Enabling checkpointing to reduce memory usage, but this slows down training. '
            'Consider enabling AMP (--amp) for fast and memory efficient training')
        torch.cuda.empty_cache()

        # 手动启用梯度检查点功能
        train_model(
            model=model,
            epochs=args.epochs,
            batch_size=args.batch_size,
            learning_rate=args.lr,
            device=device,
            img_scale=args.scale,
            val_percent=args.val / 100,
            amp=args.amp
        )
    else:
        raise e