GAN.py

原代码地址：github.com/zqhang/MTGFLOW

目录

def ConvEncoder()

def ConvDecoder()  

class CNNAE(torch.nn.Module):

class R_Net(torch.nn.Module):

class D_Net(torch.nn.Module):

def R_Loss（）

def D_Loss（）

def R_WLoss（）

def D_WLoss（）

def train_model（）

def train_single_epoch（）

def validate_single_epoch（）

def test_single_epoch（）

---

def ConvEncoder()

这个函数定义了一个卷积编码器，用于将输入数据进行特征提取。

def ConvEncoder(activation=nn.LeakyReLU, in_channels:int=3, n_c:int=64, k_size:int=5):
    """
    定义卷积编码器模型，将输入数据进行卷积和批量归一化处理。
    
    参数:
        activation: 激活函数，默认为 LeakyReLU。
        in_channels: 输入数据的通道数，默认为 3（通常是 RGB 图像的通道数）。
        n_c: 卷积核的数量，也是输出的通道数，默认为 64。
        k_size: 卷积核的大小，默认为 5。
        
    返回:
        enc: 卷积编码器模型。
    """
    # 创建一个顺序模型（Sequential Model），按顺序添加层
    enc = nn.Sequential(
        # 第一个卷积层：输入通道数为 in_channels，输出通道数为 n_c，卷积核大小为 k_size，
        # 步幅为 2，填充为 2，使用激活函数 activation
        nn.Conv1d(in_channels, n_c, k_size, stride=2, padding=2),
        # 批量归一化层，处理卷积层的输出
        nn.BatchNorm1d(n_c),
        # 激活函数
        activation(),
        
        # 第二个卷积层：输入通道数为 n_c，输出通道数为 n_c*2，卷积核大小为 k_size，
        # 步幅为 2，填充为 2，使用激活函数 activation
        nn.Conv1d(n_c, n_c*2, k_size, stride=2, padding=2),
        # 批量归一化层，处理卷积层的输出
        nn.BatchNorm1d(n_c*2),
        # 激活函数
        activation(),
        
        # 第三个卷积层：输入通道数为 n_c*2，输出通道数为 n_c*4，卷积核大小为 k_size，
        # 步幅为 2，填充为 2，使用激活函数 activation
        nn.Conv1d(n_c*2, n_c*4, k_size, stride=2, padding=2),
        # 批量归一化层，处理卷积层的输出
        nn.BatchNorm1d(n_c*4),
        # 激活函数
        activation()
    )
    # 返回卷积编码器模型
    return enc

def ConvDecoder()  

创建一个包含卷积转置层的序列模型，用于将低维特征映射回原始输入图像的高维空间。在生成对抗网络（GAN）等模型中，该函数通常用作生成器网络的一部分，负责将潜在空间（随机噪声或其他低维表示）映射为逼真的图像。卷积转置层与普通卷积层相反，它将输入扩大（上采样）而不是缩小（下采样），从而实现从低维到高维的映射。函数返回创建的卷积转置层模型。

def ConvDecoder(activation=nn.LeakyReLU, in_channels:int=3, n_c:int=64, k_size:int=5):
    # activation：激活函数，默认为 LeakyReLU
    # in_channels：输入图像的通道数，默认为3（通常是RGB图像）
    # n_c：卷积核的通道数，默认为64
    # k_size：卷积核的大小，默认为5

    # 定义一个包含卷积转置层的序列模型
    decoder = nn.Sequential(
        # 第一个卷积转置层，将输入通道数扩大4倍，然后输出通道数减半
        nn.ConvTranspose1d(n_c*4, n_c*2, k_size, stride=2, padding=2, output_padding=0),
        torch.nn.BatchNorm1d(n_c*2),  # 批归一化层，对输出进行归一化
        activation(),  # 激活函数，将输出进行非线性变换
        # 第二个卷积转置层，将输入通道数减半，然后输出通道数减半
        torch.nn.ConvTranspose1d(n_c*2, n_c, k_size, stride=2, padding=2, output_padding=1),
        torch.nn.BatchNorm1d(n_c),  # 批归一化层，对输出进行归一化
        activation(),  # 激活函数，将输出进行非线性变换
        # 第三个卷积转置层，将输入通道数减半，然后输出通道数与输入图像的通道数相同
        torch.nn.ConvTranspose1d(n_c, in_channels, k_size, stride=2, padding=2, output_padding=1)
    )
    return decoder

class CNNAE(torch.nn.Module):

这段代码定义了一个基于卷积神经网络的自动编码器模型（CNNAE）。自动编码器是一种无监督学习模型，它可以学习输入数据的紧凑表示（编码），然后再将这个紧凑表示解码为原始输入数据的重构。在这个模型中，编码器和解码器都是使用卷积神经网络实现的。模型的初始化方法（__init__）中设置了输入图像的通道数、卷积核数量和卷积核大小，并创建了编码器和解码器。前向传播方法（forward）定义了输入数据的处理过程，首先通过编码器获得特征表示，然后通过解码器将特征表示解码为重构的图像。这种结构使得模型可以学习到输入数据的有效表示，用于图像的压缩和重构等任务。

class CNNAE(torch.nn.Module):
    """自动编码器模型，使用卷积神经网络（CNN）实现"""

    def __init__(self, in_channels:int = 3, n_channels:int = 16, kernel_size:int = 5):
        # 初始化方法，定义了模型的结构和参数
        super(CNNAE, self).__init__()

        # 设置输入图像的通道数、卷积核数量、卷积核大小
        self.in_channels = in_channels
        self.n_c = n_channels
        self.k_size = kernel_size

        # 创建编码器（使用ConvEncoder类），指定激活函数和参数
        activation = torch.nn.LeakyReLU
        self.encoder = ConvEncoder(activation, in_channels, n_channels, kernel_size)

        # 创建解码器（使用ConvDecoder类），指定激活函数和参数
        self.decoder = ConvDecoder(activation, in_channels, n_channels, kernel_size)

    def forward(self, x:torch.Tensor):
        # 前向传播方法，定义了输入数据的处理过程
        # 输入x是一个张量（Tensor），代表输入图像

        # 使用编码器对输入图像进行编码，得到特征表示z
        z = self.encoder.forward(x)

        # 使用解码器对特征表示进行解码，得到重构的图像x_out
        x_out = self.decoder.forward(z)

        # 返回重构的图像x_out
        return x_out

class R_Net(torch.nn.Module):

作用： 这段代码定义了一个名为R_Net的PyTorch模型类，表示一个带有噪音的卷积自编码器。该模型包括一个Encoder和一个Decoder，Encoder用于将输入数据编码为隐藏表示，Decoder用于将隐藏表示解码为重构输出。在前向传播过程中，可以选择是否在输入中添加噪音。该模型的主要作用是学习输入数据的压缩表示，并尽可能地恢复出原始输入数据，同时能够处理带有噪音的输入。

class R_Net(torch.nn.Module):
    # 定义一个名为R_Net的PyTorch模型类，继承自torch.nn.Module基类

    def __init__(self, activation=torch.nn.LeakyReLU, in_channels:int=3, n_channels:int=16,
                 kernel_size:int=5, std:float=0.2):
        # 初始化方法，用于定义模型的结构和参数
        # 参数说明：
        # activation: 激活函数，默认为LeakyReLU
        # in_channels: 输入图像的通道数，默认为3（RGB图像）
        # n_channels: 卷积层的通道数，表示卷积核的数量，默认为16
        # kernel_size: 卷积核的大小，默认为5
        # std: 添加噪音时使用的标准差，默认为0.2

        super(R_Net, self).__init__()
        # 调用父类的构造函数，必须在子类构造函数的开始处调用

        self.activation = activation
        # 将传入的激活函数赋值给类属性self.activation

        self.in_channels = in_channels
        self.n_c = n_channels
        self.k_size = kernel_size
        self.std = std
        # 将传入的参数赋值给相应的类属性

        self.Encoder = ConvEncoder(activation, in_channels, n_channels, kernel_size)
        # 创建一个卷积编码器（使用给定的激活函数和参数）

        self.Decoder = ConvDecoder(activation, in_channels, n_channels, kernel_size)
        # 创建一个卷积解码器（使用给定的激活函数和参数）

    def forward(self, x:torch.Tensor, noise:bool=True):
        # 定义前向传播方法，定义模型的计算过程
        # 参数说明：
        # x: 输入的张量，通常是图像数据，类型为torch.Tensor
        # noise: 是否在输入中添加噪音，默认为True

        x_hat = self.add_noise(x) if noise else x
        # 如果noise为True，则在输入中添加噪音，否则不添加
        # 添加噪音的操作由self.add_noise函数完成

        z = self.Encoder.forward(x_hat)
        # 将带有噪音的输入x_hat通过Encoder模块进行编码，得到隐藏表示z

        x_out = self.Decoder.forward(z)
        # 将隐藏表示z通过Decoder模块进行解码，得到重构的输出x_out

        return x_out
        # 返回重构的输出

    def add_noise(self, x):
        # 定义一个函数，用于在输入中添加噪音
        # 参数说明：
        # x: 输入的张量，通常是图像数据，类型为torch.Tensor

        noise = torch.randn_like(x) * self.std
        # 生成与输入x相同大小的随机噪音，乘以self.std得到具有指定标准差的噪音

        x_hat = x + noise
        # 将噪音添加到输入x上，得到带有噪音的输入x_hat

        return x_hat
        # 返回带有噪音的输入x_hat

class D_Net(torch.nn.Module):

这段代码定义了一个名为D_Net的类，表示一个卷积神经网络的判别器模型。该模型接收输入图像，并输出一个单一的值，用于表示输入图像是真实样本还是生成样本。类中的forward方法定义了模型的前向传播逻辑，将输入图像通过卷积层和全连接层进行处理，得到最终的判别结果。_compute_out_dim方法用于计算卷积层输出的特征维度，以便为全连接层指定输入维度。

class D_Net(torch.nn.Module):
    def __init__(self, in_resolution:int, activation=torch.nn.LeakyReLU, in_channels:int=3, n_channels:int=16, kernel_size:int=5):
        # 初始化D_Net类，定义判别器的结构
        super(D_Net, self).__init__()
        
        # 设置激活函数和输入分辨率、通道数、卷积核大小
        self.activation = activation
        self.in_resolution = in_resolution
        self.in_channels = in_channels
        self.n_c = n_channels
        self.k_size = kernel_size

        # 创建一个卷积编码器（CNN Encoder），使用给定的激活函数和参数
        self.cnn = ConvEncoder(activation, in_channels, n_channels, kernel_size)

        # 计算D网络卷积部分输出的维度
        self.out_dim = self._compute_out_dim()

        # 创建一个全连接层，将卷积部分的输出映射到一个单一的输出值（用于二元分类）
        self.fc = torch.nn.Linear(self.out_dim, 1)
        
    def _compute_out_dim(self):
        # 计算卷积部分的输出维度，用于全连接层的输入维度
        test_x = torch.Tensor(1, self.in_channels, self.in_resolution)
        
        # 冻结卷积层的参数，防止在计算过程中被修改
        for p in self.cnn.parameters():
            p.requires_grad = False
        
        # 通过卷积部分得到输出，然后计算输出的维度
        test_x = self.cnn(test_x)
        out_dim = torch.prod(torch.tensor(test_x.shape[1:])).item()
        
        # 解冻卷积层的参数，以便在训练中更新它们
        for p in self.cnn.parameters():
            p.requires_grad = True

        return out_dim

    def forward(self, x:torch.Tensor):
        # 前向传播函数，对输入图像进行判别
        # 使用卷积层处理输入图像
        x = self.cnn(x)
        # 将卷积层输出的特征图展平成一维向量
        x = torch.flatten(x, start_dim=1)
        # 通过全连接层得到最终的判别结果
        out = self.fc(x)
        return out

def R_Loss（）

def R_Loss(d_net: torch.nn.Module, x_real: torch.Tensor, x_fake: torch.Tensor, lambd: float) -> dict:
    # d_net 是判别器模型，用于判别生成样本的真实性。
    # x_real 是真实样本的张量。
    # x_fake 是生成样本的张量。
    # lambd 是用于权衡重构损失和生成损失的权重参数。

    # pred 是生成样本经过判别器的输出，表示生成样本被判别为真实样本的概率。
    # y 是与pred相同大小的张量，其所有元素都是1，用于计算生成损失。
    pred = d_net(x_fake)
    y = torch.ones_like(pred)
    #rec_loss 是重构损失，使用均方误差（MSE）衡量生成样本x_fake与真实样本x_real之间的差异，即生成样本与真实样本的相似度。
    rec_loss = F.mse_loss(x_fake, x_real)
    # gen_loss 是生成损失，使用二元交叉熵（Binary Cross Entropy）损失函数计算生成样本被判别为真实样本的损失，即判别器预测与实际标签的差异。
    gen_loss = F.binary_cross_entropy_with_logits(pred, y) # generator loss
    # L_r 是最终的损失函数，它由生成损失和重构损失以及二者的权重参数lambd加权组成。
    L_r = gen_loss + lambd * rec_loss
    # 函数返回一个字典，包含了重构损失(rec_loss)、生成损失(gen_loss)和最终的损失(L_r)。
    # 这些损失值用于监控和优化生成模型的性能。通常，生成模型的目标是最小化生成损失，同时保持生成样本与真实样本的相似性，即最小化重构损失。
    return {'rec_loss' : rec_loss, 'gen_loss' : gen_loss, 'L_r' : L_r}

def D_Loss（）

这段代码定义了一个函数D_Loss，它计算了给定判别器（d_net）对真实样本（x_real）和生成样本（x_fake）的损失。函数首先使用判别器对真实样本和生成样本进行预测，得到预测结果。然后，为真实样本和生成样本分别创建标签（1表示真实样本，0表示生成样本）。接着，使用二元交叉熵损失函数（F.binary_cross_entropy_with_logits）分别计算真实样本和生成样本的损失（real_loss和fake_loss）。最后，将这两个损失相加，得到最终的判别器损失。该函数用于训练生成对抗网络（GAN）中的判别器，目的是使判别器能够正确区分真实样本和生成样本。

def D_Loss(d_net: torch.nn.Module, x_real: torch.Tensor, x_fake: torch.Tensor) -> torch.Tensor:
    # 输入参数:
    # d_net: 判别器模型
    # x_real: 真实样本
    # x_fake: 生成样本
    
    # 利用判别器对真实样本进行预测，得到预测结果
    pred_real = d_net(x_real)
    # 利用判别器对生成样本进行预测（使用detach()来阻止梯度回传），得到预测结果
    pred_fake = d_net(x_fake.detach())
    
    # 为真实样本和生成样本创建标签，1表示真实样本，0表示生成样本
    y_real = torch.ones_like(pred_real)
    y_fake = torch.zeros_like(pred_fake)

    # 使用二元交叉熵损失函数计算真实样本和生成样本的损失
    # real_loss表示真实样本的损失，fake_loss表示生成样本的损失
    real_loss = F.binary_cross_entropy_with_logits(pred_real, y_real)
    fake_loss = F.binary_cross_entropy_with_logits(pred_fake, y_fake)

    # 返回真实样本和生成样本损失的总和作为最终的判别器损失
    return real_loss + fake_loss

def R_WLoss（）

这段代码定义了一个函数R_WLoss，它计算了生成对抗网络（GAN）中的损失。函数使用给定的判别器（d_net）对生成样本（x_fake）进行预测，并通过 sigmoid 激活函数将预测结果映射到 [0, 1] 范围内。然后，函数计算了两个损失项：重建损失（rec_loss，使用均方误差）和生成损失（gen_loss，Wasserstein G loss）。最后，函数计算了总的损失（L_r），其中重建损失被乘以权重 lambd。函数返回一个包含损失信息的字典，其中包括重建损失、生成损失和总损失。这些损失用于优化生成器网络。

def R_WLoss(d_net: torch.nn.Module, x_real: torch.Tensor, x_fake: torch.Tensor, lambd: float) -> dict:
    # 输入参数:
    # d_net: 判别器模型
    # x_real: 真实样本
    # x_fake: 生成样本
    # lambd: 重建损失的权重
    
    # 使用判别器对生成样本进行预测，得到预测结果，并经过 sigmoid 激活函数
    pred = torch.sigmoid(d_net(x_fake))

    # 计算重建损失（均方误差）
    rec_loss = F.mse_loss(x_fake, x_real)
    # 计算生成损失（Wasserstein G loss: - E[ D(G(x)) ]）
    gen_loss = -torch.mean(pred)

    # 计算总损失（L_r = 生成损失 + 重建损失 * 权重 lambd）
    L_r = gen_loss + lambd * rec_loss

    # 返回损失信息的字典，包括重建损失、生成损失和总损失
    return {'rec_loss': rec_loss, 'gen_loss': gen_loss, 'L_r': L_r}

def D_WLoss（）

这段代码定义了一个函数D_WLoss，它计算了生成对抗网络（GAN）中鉴别器的损失。函数使用给定的判别器（d_net）对真实样本（x_real）和生成样本（x_fake）进行预测，并通过 sigmoid 激活函数将预测结果映射到 [0, 1] 范围内。然后，函数计算了两个损失项：真实样本的损失（-E[D(x_real)]）和生成样本的损失（E[D(x_fake)]）。这两个损失项相加后得到鉴别器的总损失，这个损失用于优化鉴别器网络。

def D_WLoss(d_net: torch.nn.Module, x_real: torch.Tensor, x_fake: torch.Tensor) -> torch.Tensor:
    # 输入参数:
    # d_net: 判别器模型
    # x_real: 真实样本
    # x_fake: 生成样本
    
    # 使用判别器对真实样本和生成样本进行预测，得到预测结果，并通过 sigmoid 激活函数
    pred_real = torch.sigmoid(d_net(x_real))
    pred_fake = torch.sigmoid(d_net(x_fake.detach()))
    
    # 计算鉴别器的损失（Wasserstein D loss: -E[D(x_real)] + E[D(x_fake)]）
    dis_loss = -torch.mean(pred_real) + torch.mean(pred_fake)

    # 返回鉴别器损失
    return dis_loss

def train_model（）

这段代码定义了一个函数train_model，用于训练生成对抗网络（GAN）的生成器和判别器模型。函数接受多个参数，包括生成器模型（r_net）、判别器模型（d_net）、训练数据加载器（train_loader）、测试数据加载器（test_loader）等等。函数中包含了训练循环，每个epoch会进行一次训练和测试，并根据需要保存模型。

def train_model(args, r_net: torch.nn.Module,
                d_net: torch.nn.Module,
                train_loader: torch.utils.data.DataLoader,
                test_loader: torch.utils.data.DataLoader,
                r_loss = R_Loss,
                d_loss = D_Loss,
                lr_scheduler = None,
                optimizer_class = torch.optim.Adam,
                optim_r_params: dict = {},
                optim_d_params: dict = {},
                learning_rate: float = 0.001,
                scheduler_r_params: dict = {},
                scheduler_d_params: dict = {},
                batch_size: int = 1024,
                max_epochs: int = 40,
                epoch_step: int = 1,
                save_step: int = 5,
                lambd: float = 0.2,
                device: torch.device = torch.device('cuda'),
                save_path: str = ".") -> tuple:
    # 参数说明:
    # args: 其他训练参数的配置
    # r_net: 生成器模型
    # d_net: 判别器模型
    # train_loader: 训练数据的数据加载器
    # test_loader: 测试数据的数据加载器
    # r_loss: 生成器的损失函数，默认为R_Loss
    # d_loss: 判别器的损失函数，默认为D_Loss
    # lr_scheduler: 学习率调度器，默认为None
    # optimizer_class: 优化器类型，默认为torch.optim.Adam
    # optim_r_params: 生成器优化器的参数，默认为空字典
    # optim_d_params: 判别器优化器的参数，默认为空字典
    # learning_rate: 初始学习率，默认为0.001
    # scheduler_r_params: 生成器学习率调度器的参数，默认为空字典
    # scheduler_d_params: 判别器学习率调度器的参数，默认为空字典
    # batch_size: 批大小，默认为1024
    # max_epochs: 最大训练轮数，默认为40
    # epoch_step: 每隔多少轮打印训练信息，默认为1
    # save_step: 每隔多少轮保存模型，默认为5
    # lambd: R_Loss中的lambda参数，默认为0.2
    # device: 训练设备，默认为'cuda'
    # save_path: 模型保存路径，默认为当前目录
    
    # 创建生成器和判别器的优化器
    optim_r = optimizer_class(r_net.parameters(), lr=learning_rate, **optim_r_params)
    optim_d = optimizer_class(d_net.parameters(), lr=learning_rate, **optim_d_params)

    # 创建学习率调度器
    if lr_scheduler:
        scheduler_r = lr_scheduler(optim_r, **scheduler_r_params)
        scheduler_d = lr_scheduler(optim_d, **scheduler_d_params)

    logger = log()  # 日志记录器

    # 开始训练循环
    for epoch in range(max_epochs):
        start = timer()  # 记录每轮开始时间
        # 训练一个epoch并获取训练指标
        train_metrics = train_single_epoch(r_net, d_net, optim_r, optim_d, r_loss, d_loss, train_loader, lambd, device)
        # 测试模型并获取测试指标
        gt, pre = test_metrics = test_single_epoch(r_net, d_net, r_loss, d_loss, test_loader, device)
        # 记录测试结果
        # logger.print_result(gt, pre, (r_net, d_net), args.seed, args)
        time = timer() - start  # 记录每轮训练时间

        # 每隔一定轮数打印训练信息
        if epoch % epoch_step == 0:
            print(f'Epoch {epoch}:')
            # print('Train Metrics:', train_metrics)
            # print('Test Metrics:', test_metrics)
            # print(f'TIME: {time:.2f} s')

        # 学习率调度器进行一步学习率更新
        if lr_scheduler:
            scheduler_r.step()
            scheduler_d.step()

        # 每隔一定轮数保存模型
        # if epoch % save_step == 0:
        #     torch.save(r_net.state_dict(), os.path.join(save_path, "r_net_{}.pt".format(epoch)))
        #     torch.save(d_net.state_dict(), os.path.join(save_path, "d_net_{}.pt".format(epoch)))
        #     print(f'Saving model on epoch {epoch}')

    # 返回训练好的生成器和判别器模型
    return (r_net, d_net)

def train_single_epoch（）

这段代码定义了一个函数`train_single_epoch`，用于训练一个epoch。在每个batch的训练中，首先将真实数据传入判别器，计算判别器损失，并更新判别器的参数。然后，将真实数据传入生成器生成假数据，计算生成器损失（包括重构损失和对抗损失），并更新生成器的参数。最后，将每个batch的损失累加，并计算每个样本的平均损失。函数返回一个包含平均损失的字典。

def train_single_epoch(r_net, d_net, optim_r, optim_d, r_loss, d_loss, train_loader, lambd, device) -> dict:
    # 参数说明：
    # r_net: 生成器模型
    # d_net: 判别器模型
    # optim_r: 生成器的优化器
    # optim_d: 判别器的优化器
    # r_loss: 生成器的损失函数
    # d_loss: 判别器的损失函数
    # train_loader: 训练数据的数据加载器
    # lambd: R_Loss中的lambda参数
    # device: 训练设备
    
    r_net.train()  # 设置生成器为训练模式
    d_net.train()  # 设置判别器为训练模式

    train_metrics = {'rec_loss': 0, 'gen_loss': 0, 'dis_loss': 0}  # 初始化训练指标

    for data, _, idx in train_loader:
        x = data.to(device)  # 将数据移动到指定设备上
        x = torch.transpose(x, dim0=2, dim1=3)  # 调整输入数据的维度
        x_real = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])  # 调整输入数据的形状

        x_fake = r_net(x_real)  # 通过生成器生成假数据
        d_net.zero_grad()  # 判别器梯度清零

        dis_loss = d_loss(d_net, x_real, x_fake)  # 计算判别器损失

        dis_loss.backward()  # 反向传播并更新判别器参数
        optim_d.step()  # 判别器优化器更新

        r_net.zero_grad()  # 生成器梯度清零

        r_metrics = r_loss(d_net, x_real, x_fake, lambd)  # 计算生成器损失，包含重构损失和对抗损失

        r_metrics['L_r'].backward()  # 反向传播并更新生成器参数
        optim_r.step()  # 生成器优化器更新

        # 累加每个batch的损失
        train_metrics['rec_loss'] += r_metrics['rec_loss']
        train_metrics['gen_loss'] += r_metrics['gen_loss']
        train_metrics['dis_loss'] += dis_loss

    # 计算每个样本的平均损失
    train_metrics['rec_loss'] = train_metrics['rec_loss'].item() / (len(train_loader.dataset) / train_loader.batch_size)
    train_metrics['gen_loss'] = train_metrics['gen_loss'].item() / (len(train_loader.dataset) / train_loader.batch_size)
    train_metrics['dis_loss'] = train_metrics['dis_loss'].item() / (len(train_loader.dataset) / train_loader.batch_size)

    return train_metrics  # 返回训练指标字典

def validate_single_epoch（）

这段代码定义了一个函数validate_single_epoch，用于在验证集上评估模型。与训练过程类似，首先将真实数据传入判别器，计算判别器损失。然后，将真实数据传入生成器生成假数据，计算生成器损失。最后，将每个batch的损失累加，并计算每个样本的平均损失。函数返回一个包含平均损失的字典。

def validate_single_epoch(r_net, d_net, r_loss, d_loss, valid_loader, device) -> dict:
    # 参数说明：
    # r_net: 生成器模型
    # d_net: 判别器模型
    # r_loss: 生成器的损失函数
    # d_loss: 判别器的损失函数
    # valid_loader: 验证数据的数据加载器
    # device: 训练设备

    r_net.eval()  # 设置生成器为评估模式
    d_net.eval()  # 设置判别器为评估模式

    valid_metrics = {'rec_loss': 0, 'gen_loss': 0, 'dis_loss': 0}  # 初始化验证指标

    with torch.no_grad():
        for data, _, idx in valid_loader:
            x = data.to(device)  # 将数据移动到指定设备上
            x = torch.transpose(x, dim0=2, dim1=3)  # 调整输入数据的维度
            x_real = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])  # 调整输入数据的形状

            x_fake = r_net(x_real)  # 通过生成器生成假数据

            dis_loss = d_loss(d_net, x_real, x_fake)  # 计算判别器损失

            r_metrics = r_loss(d_net, x_real, x_fake, 0)  # 计算生成器损失，lambda参数为0表示不使用重构损失

            # 累加每个batch的损失
            valid_metrics['rec_loss'] += r_metrics['rec_loss']
            valid_metrics['gen_loss'] += r_metrics['gen_loss']
            valid_metrics['dis_loss'] += dis_loss

        # 计算每个样本的平均损失
        valid_metrics['rec_loss'] = valid_metrics['rec_loss'].item() / (len(valid_loader.dataset) / valid_loader.batch_size)
        valid_metrics['gen_loss'] = valid_metrics['gen_loss'].item() / (len(valid_loader.dataset) / valid_loader.batch_size)
        valid_metrics['dis_loss'] = valid_metrics['dis_loss'].item() / (len(valid_loader.dataset) / valid_loader.batch_size)

    return valid_metrics  # 返回验证指标字典

def test_single_epoch（）

这段代码定义了一个函数test_single_epoch，用于在测试集上评估模型。与验证过程类似，首先将真实数据传入判别器，计算判别器损失，并将损失值添加到损失列表中。然后，将损失列表的元素合并为一个Tensor，并将NaN值替换为0。接着，使用roc_auc_score函数计算ROC AUC得分，并将结果打印出来。函数返回真实标签和负的损失值，用于ROC AUC计算。

def test_single_epoch(r_net, d_net, r_loss, d_loss, test_loader, device) -> dict:
    # 参数说明：
    # r_net: 生成器模型
    # d_net: 判别器模型
    # r_loss: 生成器的损失函数
    # d_loss: 判别器的损失函数
    # test_loader: 测试数据的数据加载器
    # device: 训练设备

    r_net.eval()  # 设置生成器为评估模式
    d_net.eval()  # 设置判别器为评估模式

    valid_metrics = {'rec_loss': 0, 'gen_loss': 0, 'dis_loss': 0}  # 初始化验证指标
    loss = []  # 初始化损失列表

    with torch.no_grad():
        for data, _, idx in test_loader:
            x = data.to(device)  # 将数据移动到指定设备上
            x = torch.transpose(x, dim0=2, dim1=3)  # 调整输入数据的维度
            x_real = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])  # 调整输入数据的形状

            dis_loss = d_net(x_real).squeeze().cpu()  # 计算判别器损失，并将结果移回CPU

            loss.append(dis_loss)  # 将损失添加到损失列表中

        loss = torch.cat(loss)  # 将损失列表合并为一个Tensor
        loss = np.nan_to_num(loss)  # 将NaN值替换为0
        auc_score = roc_auc_score(np.asarray(test_loader.dataset.label, dtype=int), -loss)  # 计算ROC AUC得分

        print('roc_test', auc_score)  # 打印测试集的ROC AUC得分
        return np.asarray(test_loader.dataset.label, dtype=int), -loss  # 返回真实标签和负的损失值（用于ROC AUC计算）