【论文复现记录】Image Disentanglement Autoencoder for Steganography without Embedding（IDEAS）-CVPR-2022

记录一下这篇文章的代码复现和学习过程，欢迎交流

首先要特别鸣谢我的好朋友——hxy女士，在整个复现过程中帮了我很多，她真的很优秀，感谢她百忙之中抽出时间教我，非常非常非常感谢她！！！！

一、数据集下载、处理

这篇文章用到的数据集分别是来自LSUN的卧室和教堂图像，以及来自FFHQ的人脸图像。每个子集包括70,000个随机选择的图像，归一化为256 × 256像素。下载方法有三种（参考文章：https://blog.csdn.net/qq_40859587/article/details/134563263）：

1）https://pan.baidu.com/s/17cf-7ZktjbmitZUzt9KJTg#list/path=%2F 提取码：gyl5 （开VIP速度很快）
2）http://dl.yf.io/lsun/scenes/
3）https://github.com/fyu/lsun 用脚本-c传场景参数

我选择的方法是第一种。

下载好之后是这样的：（其实没必要全下下来，当时没想到这一点，好崩溃，盘全满了）

下载完成后，发现里面的数据是以.mdb格式呈现的，需要将其转换成图片格式。先下载一个LSUN 数据集文档和演示代码，需要用到里面的data.py文件，下载链接：

fyu/lsun: LSUN Dataset Documentation and Demo Code (github.com)

然后将data.py文件放在E:\IDEAS-master下【这一步很关键】，也就是我们下载的IDEAS这篇论文的源码，在github上可以搜到。接着打开pycharm，先安装三个python依赖项：numpy、lmdb、opencv，使用如下命令：【！！注意：安装前把代理一定要关掉，否则会报错】

pip install opencv

pip install lmdb

pip install opencv-python

报错截图：

这里可以先安装一个清华源，再安装依赖项：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn opencv-python

安装完成之后就可以开始解压mdb文件了，下面是一个具体过程记录：

以classroom这个类别为例，先在F盘新建一个文件夹，用来存放转换后的图片。然后在pycharm的终端输入命令：（注意路径，我使用的是绝对路径，第一个路径是data.mdb和lock.mdb两个文件所在的位置）

python data.py export E:\LSUN\data\classroom_train_lmdb\classroom_train_lmdb --out_dir F:\classroom --flat

然后回车，可以看到已经开始转换了（转换时间好长……）

这个数据集还挺大的，将近17万张图片。

然后使用相同的方法转换classroom_val的图片：

1、F盘新建文件夹存放图片；2、pycharm终端输入命令；（这个只有300张图片，很快就转好了）

 python data.py export E:\LSUN\data\classroom_val_lmdb\classroom_val_lmdb --out_dir F:\classroom_val --flat

我一共转换了classroom、church两个类别的*_train和*_val里的图片，完成后发现图片格式为webp，需要转换成jpg格式（好崩溃，还没跑代码，到这一步已经快要歇菜了……），又查询了很多博主写的文章，选择了这篇文章的方法：python 批量webp格式转换成jpg_python图片webq转jpg格式转换-CSDN博客

真的！很好用！亲测！！！下面是我运行时的代码截图和转换成功截图：

发现图片名称乱码，使用这篇文章的教程进行批量重命名：LSUN数据集读取和解压，mdb格式转换为jpg格式（保姆教程）-CSDN博客

具体操作过程：

1、新建一个txt文件，然后输入代码：

@Echo Off&SetLocal ENABLEDELAYEDEXPANSION
FOR %%a in (*) do (
set "name=%%a"
set "name=!name: (=!"
set "name=!name:)=!"
ren "%%a" "!name!"
)
exit

然后保存，将txt改成bat格式，之后选中需要重命名的图片，右键重命名，对第一个命名为classroom_val，然后双击bat文件，即可完成。下面是重命名后的文件截图：

完成格式转换后，将数据集压缩包都上传到阿里云盘。（漫长且磨人的过程啊……）

二、数据集、代码上传

由于这篇文章的代码是在Linux上运行的，我的电脑是Windows，所以还是在AutoDL上租了一个服务器来跑：（提醒自己：服务器不用的时候一定要关机！！）

选择无卡开机（因为我们还没有正式跑代码，只是下载数据集，无卡也可以），在“快捷工具”中点击第二项“AutoPanel”，之后点击“公网网盘”，这里不用设置密码，直接点击确定，

然后使用手机端阿里云盘APP扫码授权，接着就可以下载数据集压缩包了，下载到哪里了呢，可以看到在autodl-tmp文件夹里。（继续等吧，不过这个还挺快……）

同时，还下载了一个解压工具，具体参考这篇文章：【如果服务器终端可以直接解压zip文件，就不需要这个工具，这一步可以跳过】zip压缩包太大无法用unzip成功解压（保姆级）_please check that you have transferred or created -CSDN博客

将p7zip_16.02_src_all.tar.bz2同样地，上传到阿里云盘，再下载到租的服务器中。

下载完成后，关掉autopanel，在jupyterlab的终端里解压这些包。【这一步由于没选好解压保存路径，数据集过大导致卡死，只能删了重来……】

在终端先输入cd autodl-tmp进入当前目录，然后输入解压命令：（这里的命令示例意思是将classroom_val.zip解压到classroom_val文件夹中，这个文件夹是我手动建的，为了保存不同类别的数据集图片）

unzip classroom_val.zip -d classroom_val

解压命令参考文章：linux系统中解压缩zip文件_linux能解压zip文件吗-CSDN博客

然后就可以开始训练了。

三、代码训练

使用如下命令进行训练：【各个含义在github上可以看到】

python train.py --exp_name ./experiments/1 --dataset_type normal --dataset_path ./classroom_train --num_iters 80000

这个num_iters在代码中找到，它没有设置，建议至少80k次：

我最终训练的命令为：

python train.py --exp_name ./experiments/1 --dataset_type normal --dataset_path ../classroom_train --num_iters 80000

 回车后发现缺少torchvision包：

好吧，再安装一下……（继续等……）

好的，然后每次尝试都发现有新的module没有安装，一步一步来吧……（下面是具体截图记录）

可以看到跑起来了【好开心好开心啊啊啊啊啊啊啊啊啊啊啊啊啊啊】

四、结果展示

还在跑……等跑完了再更新吧

【-----------------------------------------2024.07.06更新分割线-----------------------------------------------】

跑完的截图记录：

samples文件夹里的图片示例：（虽然看不懂……）

差不多复现结果就是这样，接下来就是学习看看具体的代码，了解含义，完成大作业报告啦。（不知道还会不会更新学习记录，谢谢大家看到这里）

-------------------------------------------【2024.07.11更新】分割线-------------------------------------------

上次只是跑了classroom_train的数据集，这次跑church_outdoor_train的【不知道是不是要分开训练，有大佬可以指点一下吗】，classroom_train跑出来的部分结果展示如下：

第一次跑代码，跑出来的结果不是很懂（不懂这四幅图分别代表什么含义），回去再看看论文吧……

然后今天在跑church_outdoor_train数据集的过程中遇到了一些报错，发现是需要重新安装ninja、lmdb、opencv_python这些库，此外还有一些简单的操作由于个人不太熟练导致报错，这里也整理出来，以备后续学习：

1、终端中切换到指定文件夹下，格式：cd ./切换相对路径

2、使用conda命令创建pytorch虚拟环境：conda create -n name python=3.x  --name为环境名称

3、跑代码前记得先激活pytorch虚拟环境，格式：conda activate ideass   --ideass是环境名称

训练代码：【注意“..”】

4、解压zip命令中要注意路径是否正确

最终训练命令如下：【注意路径是否正确】

python train.py --exp_name ./experiments/2 --dataset_type normal --dataset_path ../church_outdoor_train --num_iters 80000

 跑完了，结果展示：

------------------------------------------【2024.07.12更新】分割线------------------------------------------------

hello，跑上瘾了，打算再跑一个桥梁的数据集，方法同上：

python train.py --exp_name ./experiments/3 --dataset_type normal --dataset_path ../ibridge_train --num_iters 80000

跑完结果展示：

 

---------------------------------------------【2024.07.15更新】分割线----------------------------------------------

五、代码学习

记录代码学习过程

（一）train.py

def train(
        exp_name,  # 实验名称
        args,      # 参数
        loader,    # 数据加载器
        trainer,   # 训练器
        device     # 设备（CPU或GPU）
):
    # 将数据加载器包装在sample_data函数中，可能是为了从加载器中无限次获取样本，确保在每次迭代中都有数据供给。
    loader = sample_data(loader)
    
    # 初始化损失字典，用于存储训练过程中的损失值
    loss_dict = {}
    
    # 累积变量，用于指数移动平均的计算
    accum = 0.5 ** (32 / (10 * 1000))
    
    # 记录训练开始的时间
    start_time = time.time()
    
    # 主训练循环，从1开始迭代到args.num_iters
    for idx in range(1, args.num_iters + 1):
        # 计算当前的迭代索引
        iter_idx = idx + args.start_iter

        # 如果迭代索引超过了总迭代次数，打印"Done!"并结束训练
        if iter_idx > args.num_iters:
            print("Done!")
            break

        # 从数据加载器中获取下一个批次的样本（参考图像X）
        X = next(loader)
        
        # 将样本移动到指定的设备上（通常是GPU）
        X = X.to(device)

训练判别器部分的代码：

       '''
        Training Discriminators 训练判别器
        '''

        # 设置模型中各个模块的梯度计算状态
        # 其中，编码器（E）、生成器（G）、结构生成器（Gstru）、编码器（Ex）的梯度计算被关闭（False），以确保它们在训练判别器时不更新参数。判别器（Dreal、Dco、Ddist）的梯度计算被打开（True），以便它在训练过程中更新参数。
        requires_grad(trainer['E'], False)
        requires_grad(trainer['G'], False)
        requires_grad(trainer['Gstru'], False)
        requires_grad(trainer['Ex'], False)
        requires_grad(trainer['Dreal'], True)
        requires_grad(trainer['Dco'], True)
        requires_grad(trainer['Ddist'], True)

        # 特征编码与生成
        # 使用编码器E对输入图像X进行编码，得到结构S1和纹理T1
        S1, T1 = trainer['E'](X)
        # 随机生成一个秘密张量Z，尺寸与S1相同
        Z = torch.rand(size=(S1.shape[0], args.N, S1.shape[2], S1.shape[3]),
                       dtype=torch.float).cuda() * 2 - 1
        # 使用结构生成器Gstru对Z进行处理，生成新的结构S2
        S2 = trainer['Gstru'](Z)
        # 随机生成一个新的纹理T2，尺寸与T1相同
        T2 = torch.rand_like(T1) * 2 - 1

        # 图像合成
        # 重建图像\hat{X}_1：使用S1和T1生成
        hat_X1 = trainer['G'](S1, T1)
        # 合成图像\hat{X}_2和\hat{X}_3：分别使用S2和T1、S2和T2生成
        hat_X2 = trainer['G'](S2, T1)
        hat_X3 = trainer['G'](S2, T2)

        # 计算判别器损失L_{D,real}
        # 将生成的假图像（hat_X1, hat_X2, hat_X3）拼接并输入Dreal，得到预测结果fake_pred
        fake_pred = trainer['Dreal'](torch.cat((hat_X1, hat_X2, hat_X3), 0))
        # 将真实图像X输入Dreal，得到预测结果real_pred
        real_pred = trainer['Dreal'](X)
        # 使用对数损失计算真实和假图像的损失D_real_loss
        D_real_loss = d_logistic_loss(real_pred, fake_pred)

        # 计算判别器Dco的纹理损失L_{D,texture}
        # 将假图像hat_X2、真实图像X和参考图像X裁剪成小块
        fake_patch = patchify_image(hat_X2, args.n_crop)
        real_patch = patchify_image(X, args.n_crop)
        ref_patch = patchify_image(X, args.ref_crop * args.n_crop)
        # 将假图像块和参考图像块输入Dco，得到假图像的预测fake_texture_pred和参考输入ref_input
        fake_texture_pred, ref_input = trainer['Dco'](fake_patch, ref_patch, ref_batch=args.ref_crop)
        # 将真实图像块和参考输入再次输入Dco，得到真实图像的预测real_texture_pred
        real_texture_pred, _ = trainer['Dco'](real_patch, ref_input=ref_input)
        # 使用对数损失计算真实和假图像的纹理损失D_texture_loss
        D_texture_loss = d_logistic_loss(real_texture_pred, fake_texture_pred)

        # 计算判别器Ddist的分布损失L_{D,distribution}
        # 将假图像的纹理T1和真实图像的纹理T2分别输入Ddist，得到假图像的预测fake_dist_pred和真实图像的预测real_dist_pred
        fake_dist_pred = trainer['Ddist'](T1)
        real_dist_pred = trainer['Ddist'](T2)
        # 使用对数损失计算真实和假图像的分布损失D_dist_loss
        D_dist_loss = d_logistic_loss(real_dist_pred, fake_dist_pred)

        # 记录判别器D的损失并进行优化
        # 将各个判别器的损失记录到loss_dict中
        loss_dict["D_real_loss"] = D_real_loss
        loss_dict["D_texture_loss"] = D_texture_loss
        loss_dict["D_dist_loss"] = D_dist_loss
        # 清空判别器优化器的梯度
        trainer['d_optim'].zero_grad()
        # 反向传播计算梯度
        (D_real_loss + D_texture_loss + D_dist_loss).backward()
        # 更新判别器的参数
        trainer['d_optim'].step()

        # 正则化，每隔一定的迭代次数进行一次正则化，防止模型过拟合
        # 设置真实图像X、真实图像块real_patch和真实纹理T2的requires_grad属性为True，以便计算R1正则化损失
        if iter_idx % args.d_reg_every == 0:
            X.requires_grad = True
        # 计算真实图像的预测结果，并使用d_r1_loss函数计算R1正则化损失
            real_pred = trainer['Dreal'](X)
            D_real_r1_loss = d_r1_loss(real_pred, X)

            real_patch.requires_grad = True
            real_patch_pred, _ = trainer['Dco'](real_patch, ref_patch, ref_batch=args.ref_crop)
            D_texture_r1_loss = d_r1_loss(real_patch_pred, real_patch)

            T2.requires_grad = True
            real_uniform_pred = trainer['Ddist'](T2)
            D_dist_r1_loss = d_r1_loss(real_uniform_pred, T2)

        # 清空优化器的梯度
            trainer['d_optim'].zero_grad()
        # 计算所有R1正则化损失的总和，并进行反向传播计算梯度
            r1_loss_sum = args.real_r1 / 3 * D_real_r1_loss * args.d_reg_every
            r1_loss_sum += args.texture_r1 / 3 * D_texture_r1_loss * args.d_reg_every
            r1_loss_sum += args.dist_r1 / 3 * D_dist_r1_loss * args.d_reg_every
            r1_loss_sum.backward()
        # 使用优化器更新判别器的参数
            trainer['d_optim'].step()
        # 将R1正则化损失记录到loss_dict字典中
            loss_dict["D_real_r1_loss"] = D_real_r1_loss
            loss_dict["D_texture_r1_loss"] = D_texture_r1_loss
            loss_dict["D_dist_r1_loss"] = D_dist_r1_loss

 训练主要组件部分的代码：

       '''
        Training main components.
        '''
        # 设置模型中各个组件的梯度计算状态，启用了编码器E、生成器G、结构生成器Gstru、提取器Ex的梯度计算，禁用了判别器Dreal、Dco、Ddist的梯度计算
        requires_grad(trainer['E'], True)
        requires_grad(trainer['G'], True)
        requires_grad(trainer['Gstru'], True)
        requires_grad(trainer['Ex'], True)
        requires_grad(trainer['Dreal'], False)
        requires_grad(trainer['Dco'], False)
        requires_grad(trainer['Ddist'], False)

        # 特征编码与生成
        # 编码器E对输入图像X进行编码，得到结构特征S1和纹理特征T1
        S1, T1 = trainer['E'](X)
        # 生成随机的秘密张量Z
        Z = torch.rand(size=(S1.shape[0], args.N, S1.shape[2], S1.shape[3]),
                       dtype=torch.float).cuda() * 2 - 1
        # 使用结构生成器Gstru生成结构特征S2
        S2 = trainer['Gstru'](Z)
        # 生成与T1相同维度的随机纹理特征T2
        T2 = torch.rand_like(T1) * 2 - 1

        # 使用生成器G进行图像合成
        # 使用S1和T1合成重建图像hat_X1
        hat_X1 = trainer['G'](S1, T1)
        # 使用S2和T1合成图像hat_X2
        hat_X2 = trainer['G'](S2, T1)
        # 使用S2和T2合成图像hat_X3
        hat_X3 = trainer['G'](S2, T2)

        # 损失计算
        # L_{G,rec}
        # G_rec_loss：重建图像hat_X1和输入图像X的L1损失
        G_rec_loss = F.l1_loss(hat_X1, X)

        # L_{G,real}
        fake_pred = trainer['Dreal'](torch.cat((hat_X1, hat_X2, hat_X3), 0))
        # G_real_loss：生成图像hat_X1、hat_X2、hat_X3的非饱和生成损失
        G_real_loss = g_nonsaturating_loss(fake_pred)

        # L_{E,dist}
        fake_dist_pred = trainer['Ddist'](T1)
        # E_dist_loss：纹理特征T1的非饱和生成损失
        E_dist_loss = g_nonsaturating_loss(fake_dist_pred)

        # L_{G,texture}
        fake_patch = patchify_image(hat_X2, args.n_crop)
        ref_patch = patchify_image(X, args.ref_crop * args.n_crop)
        fake_patch_pred, _ = trainer['Dco'](fake_patch, ref_patch, ref_batch=args.ref_crop)
        # G_texture_loss：生成图像hat_X2的图像块与参考图像块的非饱和生成损失
        G_texture_loss = g_nonsaturating_loss(fake_patch_pred)

        # L_{E,stru}
        if iter_idx > args.num_iters * 0.8:
            container_image = hat_X3
        else:
            container_image = hat_X2
        # The recovered structure \hat{S}_2
        hat_S2, _ = trainer['E'](container_image)
        # E_stru_loss：从容器图像container_image中恢复的结构特征hat_S2与生成的结构特征S2的L1损失
        E_stru_loss = F.l1_loss(hat_S2, S2)

        # L_{REC}
        # The extracted secret tensor \hat{Z}
        hat_Z = trainer['Ex'](hat_S2)
        # Ex_loss：从hat_S2中提取的秘密张量hat_Z与生成的秘密张量Z的L1损失
        Ex_loss = F.l1_loss(hat_Z, Z)

        # 损失记录
        # 将各个损失记录到loss_dict字典中
        loss_dict["G_rec_loss"] = G_rec_loss
        loss_dict["G_real_loss"] = G_real_loss
        loss_dict["G_texture_loss"] = G_texture_loss
        loss_dict["E_dist_loss"] = E_dist_loss
        loss_dict["E_stru_loss"] = E_stru_loss
        loss_dict["Ex_loss"] = Ex_loss

        # 计算生成器和编码器的总损失Loss_G、Loss_E和提取器的损失Loss_Ex
        # L_G
        Loss_G = G_rec_loss + G_texture_loss + 2 * G_real_loss
        # L_E
        Loss_E = E_dist_loss + E_stru_loss
        # L_Ex
        Loss_Ex = Ex_loss
        # L_{total}:计算总损失Loss_total
        Loss_total = Loss_G + Loss_E + args.lambda_Ex * Loss_Ex

        # 优化生成器和编码器的参数
        trainer['g_optim'].zero_grad()
        Loss_total.backward(retain_graph=True)
        trainer['g_optim'].step()

        # 优化提取器的参数
        trainer['ex_optim'].zero_grad()
        Loss_Ex.backward()
        trainer['ex_optim'].step()

        # EMA更新：使用指数移动平均(EMA)更新模型参数E_ema、G_ema、Gstru_ema、Ex_ema
        accumulate(trainer['E_ema'], trainer['E'], accum)
        accumulate(trainer['G_ema'], trainer['G'], accum)
        accumulate(trainer['Gstru_ema'], trainer['Gstru'], accum)
        accumulate(trainer['Ex_ema'], trainer['Ex'], accum)

        # Log日志记录（每隔一定的迭代次数记录一次损失值和训练时间，并打印和保存日志）
        if iter_idx % args.log_every == 0:
            G_rec_val = loss_dict["G_rec_loss"].mean().item()
            G_texture_val = loss_dict["G_texture_loss"].mean().item()
            G_real_val = loss_dict["G_real_loss"].mean().item()

            E_dist_val = loss_dict["E_dist_loss"].mean().item()
            E_stru_val = loss_dict["E_stru_loss"].mean().item()

            Ex_val = loss_dict["Ex_loss"].mean().item()

            now_time = time.time()
            used_time = now_time - start_time
            rest_time = (now_time - start_time) / idx * (args.num_iters - iter_idx)

            log_output = f"[{iter_idx:07d}/{args.num_iters:07}] Total: {Loss_total.item():.4f}; " \
                         f"G,rec: {G_rec_val:.4f}; G,texture: {G_texture_val:.4f}; G,real: {G_real_val:.4f}; " \
                         f"E,dist: {E_dist_val:.4f}; E,stru: {E_stru_val:.4f}; Ex: {Ex_val:.4f} " \
                         f"used time: {time_change(used_time)};" \
                         f"rest time: {time_change(rest_time)}"

            print(log_output, flush=True)
            with open(f'{base_dir}/training_logs.txt', 'a') as fp:
                fp.write(f'{log_output}\n')

        # 输出样本：输出训练过程中生成的样本图像，并记录相关日志
        # 每隔一定的迭代次数，采样生成图像并保存到指定目录
        if iter_idx % args.show_every == 0:
            with torch.no_grad():
                # Sample a secret message and map it to secret tensor
                S1, T1 = trainer['E_ema'](X)
                # The secret message M
                M = torch.randint(low=0, high=2, dtype=torch.float,
                                  size=(S1.shape[0], args.N * S1.shape[2] * S1.shape[3]))
                Z = message_to_tensor(M, sigma=1, delta=0.5).to(device)
                Z = Z.reshape(shape=(S1.shape[0], args.N, S1.shape[2], S1.shape[3]))

                # Generate structure S2 from the secret tensor
                S2 = trainer['Gstru_ema'](Z)

                # Sample a texture T2
                T2 = torch.rand_like(T1) * 2 - 1

                # Image Synthesis
                hat_X1 = trainer['G_ema'](S1, T1)
                hat_X2 = trainer['G_ema'](S2, T1)
                hat_X3 = trainer['G_ema'](S2, T2)

                # Secret Tensor Extracting
                if iter_idx > args.num_iters * 0.8:
                    container_image = hat_X3
                    fake_img_used_as_container = 3
                else:
                    container_image = hat_X2
                    fake_img_used_as_container = 2
                hat_S2, _ = trainer['E_ema'](container_image)
                hat_Z = trainer['Ex_ema'](hat_S2)

                tensor_recovering_loss = torch.mean(torch.abs(hat_Z - Z))
                hat_Z = hat_Z.reshape(shape=(Z.shape[0], -1))
                # The extracted secret message \hat_{M}
                hat_M = tensor_to_message(hat_Z, sigma=1)

                BER = torch.mean(torch.abs(M - hat_M))
                ACC = 1 - BER

                log_output = f'[Testing {iter_idx:07d}/{args.num_iters:07d}] sigma=1 delta=50% ' \
                             f'using synthesised image \hatX_{fake_img_used_as_container} ' \
                             f'ACC of Msg: {ACC:.4f}; L1 loss of tensor: {tensor_recovering_loss:.4f}'
                print(log_output, flush=True)
                with open(f'{base_dir}/training_logs.txt', 'a') as fp:
                    fp.write(f'{log_output}\n')

                sample = torch.cat((X, hat_X1, hat_X2, hat_X3), 0)

                utils.save_image(
                    sample,
                    f"{sample_dir}/{iter_idx:07d}.png",
                    nrow=int(args.batch_size),
                    normalize=True,
                    range=(-1, 1),
                )

                print(f'Sample images are saved in experiments/{exp_name}/samples')

        # 模型保存：每隔一定的迭代次数，将模型参数和训练状态保存到指定的检查点文件中
        if iter_idx % args.save_every == 0:
            trainer_ckpt = {}
            for key in trainer.keys():
                trainer_ckpt[key] = trainer[key].state_dict()
            torch.save(
                {
                    'iter_idx': iter_idx,
                    'N': args.N,
                    "trainer": trainer_ckpt,
                    "args": args,
                },
                f"{ckpt_dir}/{iter_idx}.pt",
            )

            print(f'Checkpoint is saved in experiments/{exp_name}/checkpoints')

整个训练流程的初始化和配置部分：（定义了训练过程中的各种参数，并初始化了所需的模型、优化器和数据加载器。）

if __name__ == "__main__":
    # 设置设备和CUDA优化
    device = "cuda"
    torch.backends.cudnn.benchmark = True

    # 解析命令行参数，通过命令行参数来配置实验名称、数据集路径、训练迭代次数、学习率、批次大小等超参数
    parser = argparse.ArgumentParser()
    # Working directory: experiments/exp_name
    parser.add_argument("--exp_name", type=str, required=True)
    # Training dataset
    parser.add_argument("--dataset_path", type=str, required=True)
    # Select 'lmdb' for the lmdb files, like LSUN (https://github.com/fyu/lsun)
    # Select 'normal' for the dataset storing files (e.g., in PNG format) in a folder, like FFHQ (https://github.com/NVlabs/ffhq-dataset)
    parser.add_argument("--dataset_type", choices=['lmdb', 'normal'], required=True)

    # We recommend training at least 80k iterations
    parser.add_argument("--num_iters", type=int, required=True)
    # Hyper-parameters
    parser.add_argument("--N", type=int, default=1)
    parser.add_argument("--lambda_Ex", type=float, default=10)
    # Resume training
    parser.add_argument("--ckpt", type=str, default=None)

    # Trainig parameters
    parser.add_argument("--lr", type=float, default=0.002)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--image_size", type=int, default=256)
    parser.add_argument("--real_r1", type=float, default=10)
    parser.add_argument("--texture_r1", type=float, default=1)
    parser.add_argument("--dist_r1", type=float, default=1)
    parser.add_argument("--ref_crop", type=int, default=4)
    parser.add_argument("--n_crop", type=int, default=8)
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--channel", type=int, default=32)
    parser.add_argument("--channel_multiplier", type=int, default=1)
    parser.add_argument("--structure_channel", type=int, default=8)
    parser.add_argument("--texture_channel", type=int, default=2048)

    # Output logs every 'log_every' iterations
    parser.add_argument("--log_every", type=int, default=200)
    # Save example images every 'show_every' iterations
    parser.add_argument("--show_every", type=int, default=1000)
    # Save models every 'save_every' iterations
    parser.add_argument("--save_every", type=int, default=200000)

    args = parser.parse_args()
    args.start_iter = 0
    args.blur_kernel = (1, 3, 3, 1)

    # 创建工作目录和日志文件
    # 创建实验所需的目录结构，包括检查点和样本保存目录
    base_dir = f"experiments/{args.exp_name}"
    ckpt_dir = f"{base_dir}/checkpoints"
    sample_dir = f"{base_dir}/samples"
    os.makedirs(ckpt_dir, exist_ok=True)
    os.makedirs(sample_dir, exist_ok=True)

    # 保存训练配置到文本文件中
    with open(f"{base_dir}/training_config.txt", "wt") as fp:
        for k, v in vars(args).items():
            fp.write(f'{k}: {v}\n')
        fp.close()

    # 清空之前的训练日志文件
    with open(f"{base_dir}/training_logs.txt", "wt") as fp:
        fp.close()

    # 初始化模型（Init models）：初始化训练所需的模型，包括编码器、生成器、结构生成器、提取器、以及多个判别器，同时初始化这些模型的EMA（指数移动平均）版本。
    trainer = {
        'E': init_model('DisentanglementEncoder', args).to(device),
        'G': init_model('Generator', args).to(device),
        'Gstru': init_model('StructureGenerator', args).to(device),
        'Ex': init_model('TensorExtractor', args).to(device),

        'Dreal': init_model('ImageLevelDiscriminator', args).to(device),
        'Dco': init_model('CooccurenceDiscriminator', args).to(device),
        'Ddist': init_model('DistributionDiscriminator', args).to(device),

        'E_ema': init_model('DisentanglementEncoder', args).to(device),
        'G_ema': init_model('Generator', args).to(device),
        'Gstru_ema': init_model('StructureGenerator', args).to(device),
        'Ex_ema': init_model('TensorExtractor', args).to(device),
    }
    # 将EMA模型设置为评估模式，并用初始模型参数更新EMA模型
    trainer['E_ema'].eval()
    trainer['G_ema'].eval()
    trainer['Gstru_ema'].eval()
    trainer['Ex_ema'].eval()

    accumulate(trainer['E_ema'], trainer['E'], 0)
    accumulate(trainer['G_ema'], trainer['G'], 0)
    accumulate(trainer['Gstru_ema'], trainer['Gstru'], 0)
    accumulate(trainer['Ex_ema'], trainer['Ex'], 0)

    # 初始化生成器、提取器和判别器的优化器，分别设置不同的学习率和betas参数
    trainer['g_optim'] = optim.Adam(
        list(trainer['E'].parameters()) + list(trainer['G'].parameters()) + list(trainer['Gstru'].parameters()),
        lr=args.lr,
        betas=(0, 0.99),
    )
    trainer['ex_optim'] = optim.Adam(
        trainer['Ex'].parameters(),
        lr=args.lr,
        betas=(0, 0.99),
    )
    d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
    trainer['d_optim'] = optim.Adam(
        list(trainer['Dreal'].parameters()) + list(trainer['Dco'].parameters()) + list(trainer['Ddist'].parameters()),
        lr=args.lr * d_reg_ratio,
        betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
    )

    # 加载检查点（可选）：如果指定了检查点文件，则加载该文件中的模型参数和训练状态
    # Resume training from the 'experiments/exp_name/checkpoints/{ckpt}.pt' file
    if args.ckpt is not None:
        print("load model:", args.ckpt)
        ckpt = torch.load(f"{ckpt_dir}/{args.ckpt}.pt", map_location=lambda storage, loc: storage)
        args.start_iter = ckpt['iter_idx']
        for key in trainer.keys():
            trainer[key].load_state_dict(ckpt['trainer'][key])
    else:
        args.start_iter = 0

    # Init transforms：初始化数据增强变换（如随机水平翻转、归一化等）
    transform = transforms.Compose(
        [
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )

    # Init dataset and dataloader初始化数据集及其加载器
    dataset = set_dataset(
        type=args.dataset_type,
        path=args.dataset_path,
        transform=transform,
        resolution=args.image_size
    )
    loader = data.DataLoader(
        dataset=dataset,
        batch_size=args.batch_size,
        sampler=data_sampler(dataset=dataset, shuffle=True)
    )

    # 启动训练过程
    print('Data Loaded')
    exp_name = args.exp_name
    # 调用train函数开始训练模型，传入实验名称、参数、数据加载器、训练器和设备
    train(
        exp_name=exp_name,
        args=args,
        loader=loader,
        trainer=trainer,
        device=device
    )

（二）models.py

1、EqualConvTranspose2d 类

EqualConvTranspose2d 是一个卷积转置层，使用等效学习率技术来稳定训练。

class EqualConvTranspose2d(nn.Module):
    def __init__(
            self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
    ):
        super().__init__()

        self.weight = nn.Parameter(
            torch.randn(in_channel, out_channel, kernel_size, kernel_size)
        )
        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)

        self.stride = stride
        self.padding = padding

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channel))

        else:
            self.bias = None

    def forward(self, input):
        out = F.conv_transpose2d(
            input,
            self.weight * self.scale,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
        )

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.weight.shape[0]}, {self.weight.shape[1]},"
            f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
        )

• 初始化：

  • self.weight：卷积核权重，初始化为高斯分布。
  • self.scale：缩放因子，用于等效学习率，确保权重的方差保持稳定。
  • self.stride 和 self.padding：卷积操作的步长和填充。
  • self.bias：偏置项，可选，如果有偏置则初始化为零。

• 前向传播：

  • F.conv_transpose2d：执行转置卷积操作，使用缩放后的权重和偏置。
  • self.weight * self.scale：调整后的权重。

• __repr__ 方法：

  • 返回层的字符串表示，便于打印和调试。

2、ConvLayer 类

ConvLayer 是一个顺序卷积层，可以执行上采样、下采样和不同的填充方式。

class ConvLayer(nn.Sequential):
    def __init__(
            self,
            in_channel,
            out_channel,
            kernel_size,
            upsample=False,
            downsample=False,
            blur_kernel=(1, 3, 3, 1),
            bias=True,
            activate=True,
            padding="zero",
            tanh=False
    ):
        layers = []

        self.padding = 0
        stride = 1

        if downsample:
            factor = 2
            p = (len(blur_kernel) - factor) + (kernel_size - 1)
            pad0 = (p + 1) // 2
            pad1 = p // 2

            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))

            stride = 2

        if upsample:
            layers.append(
                EqualConvTranspose2d(
                    in_channel,
                    out_channel,
                    kernel_size,
                    padding=0,
                    stride=2,
                    bias=bias and not activate,
                )
            )

            factor = 2
            p = (len(blur_kernel) - factor) - (kernel_size - 1)
            pad0 = (p + 1) // 2 + factor - 1
            pad1 = p // 2 + 1

            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))

        else:
            if not downsample:
                if padding == "zero":
                    self.padding = (kernel_size - 1) // 2

                elif padding == "reflect":
                    padding = (kernel_size - 1) // 2

                    if padding > 0:
                        layers.append(nn.ReflectionPad2d(padding))

                    self.padding = 0

                elif padding != "valid":
                    raise ValueError('Padding should be "zero", "reflect", or "valid"')

            layers.append(
                EqualConv2d(
                    in_channel,
                    out_channel,
                    kernel_size,
                    padding=self.padding,
                    stride=stride,
                    bias=bias and not activate,
                )
            )

        if activate:
            if tanh:
                layers.append(nn.Tanh())
            else:
                if bias:
                    layers.append(FusedLeakyReLU(out_channel))

                else:
                    layers.append(ScaledLeakyReLU(0.2))

        super().__init__(*layers)

• 初始化：

  • 根据上采样、下采样和填充选项动态构建层。
  • blur_kernel：模糊核，用于在下采样和上采样时减少伪影。
  • stride 和 padding：控制卷积的步长和填充。
  • activate：决定是否添加激活函数。

• 上采样：

  • 使用 EqualConvTranspose2d 进行上采样。
  • 添加模糊层 Blur 以减少伪影。

• 下采样：

  • 使用模糊层 Blur 进行预处理。
  • 设置步长为 2 以实现下采样。

• 填充：

  • 支持三种填充方式：“zero”、“reflect”和“valid”。

• 激活函数：

  • 支持两种激活函数：FusedLeakyReLU 和 ScaledLeakyReLU。
  • 还支持 tanh 激活函数。

3、StyledResBlock 类

StyledResBlock 是一个带有样式输入的残差块，通常用于生成对抗网络（GAN）或图像风格转换任务。

class StyledResBlock(nn.Module):
    def __init__(
            self, in_channel, out_channel, style_dim, upsample, blur_kernel=(1, 3, 3, 1)
    ):
        super().__init__()

        self.conv1 = StyledConv(
            in_channel,
            out_channel,
            3,
            style_dim,
            upsample=upsample,
            blur_kernel=blur_kernel,
        )

        self.conv2 = StyledConv(out_channel, out_channel, 3, style_dim)

        if upsample or in_channel != out_channel:
            self.skip = ConvLayer(
                in_channel,
                out_channel,
                1,
                upsample=upsample,
                blur_kernel=blur_kernel,
                bias=False,
                activate=False,
            )

        else:
            self.skip = None

    def forward(self, input, style, noise=None):
        out = self.conv1(input, style, noise)
        out = self.conv2(out, style, noise)

        if self.skip is not None:
            skip = self.skip(input)

        else:
            skip = input

        return (out + skip) / math.sqrt(2)

• 初始化：

  • conv1 和 conv2：两个 StyledConv 层，用于带有样式输入的卷积操作。
  • skip：用于跳跃连接，如果需要上采样或输入输出通道数不同，则添加一个 ConvLayer 层。

• 前向传播：

  • conv1 和 conv2：依次执行两个卷积操作。
  • skip：执行跳跃连接。
  • 返回加权平均后的结果，保留残差连接的优点。

4、ResBlock 类

ResBlock 是一个标准的残差块，常用于各种卷积神经网络。

class ResBlock(nn.Module):
    def __init__(
            self,
            in_channel,
            out_channel,
            downsample,
            padding="zero",
            blur_kernel=(1, 3, 3, 1),
    ):
        super().__init__()

        self.conv1 = ConvLayer(in_channel, out_channel, 3, padding=padding)

        self.conv2 = ConvLayer(
            out_channel,
            out_channel,
            3,
            downsample=downsample,
            padding=padding,
            blur_kernel=blur_kernel,
        )

        if downsample or in_channel != out_channel:
            self.skip = ConvLayer(
                in_channel,
                out_channel,
                1,
                downsample=downsample,
                blur_kernel=blur_kernel,
                bias=False,
                activate=False,
            )

        else:
            self.skip = None

    def forward(self, input):
        out = self.conv1(input)
        out = self.conv2(out)

        if self.skip is not None:
            skip = self.skip(input)

        else:
            skip = input

        return (out + skip) / math.sqrt(2)

• 初始化：

  • conv1 和 conv2：两个 ConvLayer 层，执行卷积操作。
  • skip：用于跳跃连接，如果需要下采样或输入输出通道数不同，则添加一个 ConvLayer 层。

• 前向传播：

  • conv1 和 conv2：依次执行两个卷积操作。
  • skip：执行跳跃连接。
  • 返回加权平均后的结果，保留残差连接的优点。

5、DisentanglementEncoder 类

DisentanglementEncoder 是一个用于图像特征分离的编码器，将图像分离成结构特征和纹理特征。

class DisentanglementEncoder(nn.Module):
    def __init__(
            self,
            channel,
            structure_channel=8,
            texture_channel=2048,
            blur_kernel=(1, 3, 3, 1),
    ):
        super().__init__()

        stem = [ConvLayer(3, channel, 1)]

        in_channel = channel
        for i in range(1, 5):
            ch = channel * (2 ** i)
            stem.append(ResBlock(in_channel, ch, downsample=True, padding="reflect", blur_kernel=blur_kernel))
            in_channel = ch

        self.stem = nn.Sequential(*stem)

        self.structure = nn.Sequential(
            ConvLayer(in_channel, in_channel, 1, blur_kernel=blur_kernel),
            ConvLayer(in_channel, structure_channel, 1, blur_kernel=blur_kernel)
        )

        self.texture = nn.Sequential(
            ConvLayer(in_channel, in_channel * 2, 3, downsample=True, padding="valid", blur_kernel=blur_kernel),
            ConvLayer(in_channel * 2, in_channel * 4, 3, downsample=True, padding="valid", blur_kernel=blur_kernel),
            nn.AdaptiveAvgPool2d(1),
            ConvLayer(in_channel * 4, texture_channel, 1, tanh=True, blur_kernel=blur_kernel),
        )

    def forward(self, input):
        out = self.stem(input)

        structure = self.structure(out)
        texture = torch.flatten(self.texture(out), 1)

        return structure, texture

• 初始化：

  • stem：由多个 ResBlock 组成的主干网络，逐步下采样输入图像。
  • structure：用于提取结构特征的网络，由两个 ConvLayer 层组成。
  • texture：用于提取纹理特征的网络，包括下采样卷积层和自适应平均池化层，最后输出一个 tanh 激活的卷积层。

• 前向传播：

  • stem：通过主干网络提取初步特征。
  • structure：提取结构特征。
  • texture：提取纹理特征，并展平为一维向量。

6、Generator 类

Generator 类用于生成图像，其输入是结构特征和纹理特征，输出是生成的图像。

class Generator(nn.Module):
    def __init__(
            self,
            channel,
            structure_channel=8,
            texture_channel=2048,
            blur_kernel=(1, 3, 3, 1),
    ):
        super().__init__()

        ch_multiplier = (4, 8, 12, 16, 16, 16, 8, 4)
        upsample = (False, False, False, False, True, True, True, True)

        self.layers = nn.ModuleList()
        in_ch = structure_channel
        for ch_mul, up in zip(ch_multiplier, upsample):
            self.layers.append(
                StyledResBlock(
                    in_ch, channel * ch_mul, texture_channel, up, blur_kernel
                )
            )
            in_ch = channel * ch_mul

        self.to_rgb = ConvLayer(in_ch, 3, 1, activate=False)

    def forward(self, structure, texture, noises=None):
        if noises is None:
            noises = [None] * len(self.layers)

        out = structure
        for layer, noise in zip(self.layers, noises):
            out = layer(out, texture, noise)

        out = self.to_rgb(out)

        return out

• 初始化：

  • ch_multiplier 和 upsample：分别定义了每层的通道数倍数和是否上采样。
  • self.layers：包含多个 StyledResBlock 层，每层可以选择是否进行上采样。
  • to_rgb：将最后一层的输出转化为 RGB 图像。

• 前向传播：

  • 初始化 noises：如果没有提供噪声，则初始化为 None。
  • 依次通过每个 StyledResBlock 层。
  • 最后通过 to_rgb 层将特征转换为 RGB 图像。

7、StructureGenerator 类

StructureGenerator 类用于生成结构特征，其输入是噪声，输出是结构特征。

class StructureGenerator(nn.Module):
    def __init__(
            self,
            channel,
            N=1,
            structure_channel=8,
            blur_kernel=(1, 3, 3, 1),
    ):
        super().__init__()

        stem = [ConvLayer(N, channel, 1, blur_kernel=blur_kernel)]
        stem.append(ResBlock(channel, channel * 2, downsample=False, padding="reflect", blur_kernel=blur_kernel))
        stem.append(ResBlock(channel * 2, channel * 4, downsample=False, padding="reflect", blur_kernel=blur_kernel))
        stem.append(ResBlock(channel * 4, channel * 2, downsample=False, padding="reflect", blur_kernel=blur_kernel))
        stem.append(ConvLayer(channel * 2, structure_channel, 1, blur_kernel=blur_kernel))

        self.structure = nn.Sequential(*stem)

    def forward(self, noise):
        out = self.structure(noise)
        return out

• 初始化：

  • stem：包含多个 ConvLayer 和 ResBlock 层，用于逐步提取和处理特征。

• 前向传播：

  • 输入噪声通过 self.structure 层，逐步提取和处理，输出结构特征。

8、ImageLevelDiscriminator 类

ImageLevelDiscriminator 类用于判别生成图像的真实性，其输入是图像，输出是真实性评分。

class ImageLevelDiscriminator(nn.Module):
    def __init__(self, size, channel_multiplier=1, blur_kernel=(1, 3, 3, 1)):
        super().__init__()

        channels = {
            4: 512,
            8: 512,
            16: 512,
            32: 512,
            64: 256 * channel_multiplier,
            128: 128 * channel_multiplier,
            256: 64 * channel_multiplier,
            512: 32 * channel_multiplier,
            1024: 16 * channel_multiplier,
        }

        convs = [ConvLayer(3, channels[size], 1, blur_kernel=blur_kernel)]

        log_size = int(math.log(size, 2))

        in_channel = channels[size]

        for i in range(log_size, 2, -1):
            out_channel = channels[2 ** (i - 1)]

            convs.append(ResBlock(in_channel, out_channel, downsample=True, blur_kernel=blur_kernel))

            in_channel = out_channel

        self.convs = nn.Sequential(*convs)

        self.final_conv = ConvLayer(in_channel, channels[4], 3, blur_kernel=blur_kernel)
        self.final_linear = nn.Sequential(
            EqualLinear(channels[4] * 4 * 4, channels[4], activation="fused_lrelu"),
            EqualLinear(channels[4], 1),
        )

    def forward(self, input):
        out = self.convs(input)
        out = self.final_conv(out)

        out = out.view(out.shape[0], -1)
        out = self.final_linear(out)

        return out

• 初始化：

  • channels：定义了每个分辨率对应的通道数。
  • convs：包含多个 ConvLayer 和 ResBlock 层，逐步下采样和提取特征。
  • final_conv 和 final_linear：最终卷积层和全连接层，用于输出真实性评分。

• 前向传播：

  • 输入图像通过 self.convs 层，逐步下采样和提取特征。
  • 最后通过 final_conv 和 final_linear 层输出真实性评分。

 9、CooccurenceDiscriminator 类

CooccurenceDiscriminator 类用于比较输入图像和参考图像之间的特征相似性，输出一个判别结果。

class CooccurenceDiscriminator(nn.Module):
    def __init__(self, channel, size=256):
        super().__init__()

        encoder = [ConvLayer(3, channel, 1)]

        ch_multiplier = (2, 4, 8, 12, 12, 24)
        downsample = (True, True, True, True, True, False)
        in_ch = channel
        for ch_mul, down in zip(ch_multiplier, downsample):
            encoder.append(ResBlock(in_ch, channel * ch_mul, down))
            in_ch = channel * ch_mul

        if size > 511:
            k_size = 3
            feat_size = 2 * 2

        else:
            k_size = 2
            feat_size = 1 * 1

        encoder.append(ConvLayer(in_ch, channel * 12, k_size, padding="valid"))

        self.encoder = nn.Sequential(*encoder)

        self.linear = nn.Sequential(
            EqualLinear(
                channel * 12 * 2 * feat_size, channel * 32, activation="fused_lrelu"
            ),
            EqualLinear(channel * 32, channel * 32, activation="fused_lrelu"),
            EqualLinear(channel * 32, channel * 16, activation="fused_lrelu"),
            EqualLinear(channel * 16, 1),
        )

    def forward(self, input, reference=None, ref_batch=None, ref_input=None):
        out_input = self.encoder(input)

        if ref_input is None:
            ref_input = self.encoder(reference)
            _, channel, height, width = ref_input.shape
            ref_input = ref_input.view(-1, ref_batch, channel, height, width)
            ref_input = ref_input.mean(1)

        out = torch.cat((out_input, ref_input), 1)
        out = torch.flatten(out, 1)
        out = self.linear(out)

        return out, ref_input

• 初始化：

  • encoder：包含多个 ConvLayer 和 ResBlock 层，用于逐步提取和处理特征。
  • linear：全连接层序列，用于输出最终的判别结果。

• 前向传播：

  • input 通过 encoder 层得到 out_input。
  • 如果 ref_input 为空，则将 reference 图像也通过 encoder 层处理，并计算其平均值。
  • 将 out_input 和 ref_input 连接起来，展平后通过 linear 层得到最终的判别结果 out。

10、DistributionDiscriminator 类

DistributionDiscriminator 类用于对纹理特征进行分布判别。

class DistributionDiscriminator(nn.Module):
    def __init__(self, texture_channel=2048):
        super().__init__()
        self.model = nn.Sequential(
            EqualLinear(texture_channel, texture_channel // 4, activation="fused_lrelu"),
            EqualLinear(texture_channel // 4, texture_channel // 16, activation="fused_lrelu"),
            EqualLinear(texture_channel // 16, texture_channel // 64, activation="fused_lrelu"),
            EqualLinear(texture_channel // 64, 1, activation="fused_lrelu")
        )

    def forward(self, input):
        out = self.model(input)
        return out

• 初始化：

  • model：包含多个全连接层，用于对纹理特征进行处理和判别。

• 前向传播：

  • 输入 input 通过 model 层得到输出 out，表示纹理特征的判别结果。

11、TensorExtractor 类

TensorExtractor 类用于从结构特征中提取张量。

class TensorExtractor(nn.Module):
    def __init__(
            self,
            channel,
            N=1,
            structure_channel=8,
            blur_kernel=(1, 3, 3, 1),
    ):
        super().__init__()

        stem = [ConvLayer(structure_channel, channel * 2, 1, blur_kernel=blur_kernel)]
        stem.append(ResBlock(channel * 2, channel * 4, downsample=False, padding="reflect", blur_kernel=blur_kernel))
        stem.append(ResBlock(channel * 4, channel * 2, downsample=False, padding="reflect", blur_kernel=blur_kernel))
        stem.append(ResBlock(channel * 2, channel, downsample=False, padding="reflect", blur_kernel=blur_kernel))
        stem.append(ConvLayer(channel, N, 1, blur_kernel=blur_kernel))

        self.extract = nn.Sequential(*stem)

    def forward(self, input):
        out = self.extract(input)

        return out

• 初始化：

  • stem：包含多个 ConvLayer 和 ResBlock 层，用于逐步提取和处理特征。

• 前向传播：

  • 输入 input 通过 extract 层逐步提取特征，最终输出 out。

12、init_model 函数

根据传入的 model 字符串和 args 参数初始化并返回相应的模型对象。

• 参数说明：
  • model：字符串，指定要初始化的模型类型。
  • args：包含模型初始化所需的参数的命名空间（可能是 argparse.Namespace 对象或类似的结构）。

def init_model(model, args):
    if model == 'DisentanglementEncoder':
        return DisentanglementEncoder(
            channel=args.channel,
            structure_channel=args.structure_channel,
            texture_channel=args.texture_channel,
            blur_kernel=args.blur_kernel
        )
    elif model == 'Generator':
        return Generator(
            channel=args.channel,
            structure_channel=args.structure_channel,
            texture_channel=args.texture_channel,
            blur_kernel=args.blur_kernel
        )
    elif model == 'StructureGenerator':
        return StructureGenerator(
            channel=args.channel,
            N=args.N,
            structure_channel=args.structure_channel,
            blur_kernel=args.blur_kernel
        )
    elif model == 'ImageLevelDiscriminator':
        return ImageLevelDiscriminator(
            size=args.image_size,
            channel_multiplier=args.channel_multiplier,
            blur_kernel=args.blur_kernel
        )
    elif model == 'CooccurenceDiscriminator':
        return CooccurenceDiscriminator(
            channel=args.channel,
            size=args.image_size
        )
    elif model == 'DistributionDiscriminator':
        return DistributionDiscriminator(
            texture_channel=args.texture_channel
        )
    elif model == 'TensorExtractor':
        return TensorExtractor(
            channel=args.channel,
            N=args.N,
            structure_channel=args.structure_channel,
            blur_kernel=args.blur_kernel
        )
    else:
        raise NotImplementedError

• 分支解析：

  • 根据 model 字符串的不同值，选择不同的模型进行初始化并返回。
  • 每个分支根据对应的模型类型创建相应的实例，并传入相应的参数 args。

• 模型初始化：

  • 每个分支都调用对应模型类的构造函数，并根据 args 中的具体参数来初始化模型。
  • 参数包括但不限于 channel（通道数）、structure_channel（结构通道数）、texture_channel（纹理通道数）、blur_kernel（模糊核大小）、size（图像尺寸）、channel_multiplier（通道倍增因子）等。

• 异常处理：

  • 如果传入的 model 不在预定义的模型类型中，则抛出 NotImplementedError 异常。

（三）dataset.py

        定义了两个自定义的 PyTorch 数据集类 LMDBDataset 和 NormalDataset，以及一个工具函数 set_dataset 来根据指定的数据集类型创建相应的数据集实例。

1、LMDBDataset 类

 功能是从 LMDB 数据库中加载图像数据。

class LMDBDataset(Dataset):
    def __init__(self, path, transform, resolution=256, max_num=70000):
        self.env = lmdb.open(
            path,
            max_readers=32,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )

        if not self.env:
            raise IOError('Cannot open lmdb dataset', path)

        self.keys = []
        with self.env.begin(write=False) as txn:
            cursor = txn.cursor()
            for idx, (key, _) in enumerate(cursor):
                self.keys.append(key)
                if idx > max_num:
                    break

        self.length = len(self.keys)
        self.resolution = resolution
        self.transform = transform

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        with self.env.begin(write=False) as txn:
            key = self.keys[index]
            img_bytes = txn.get(key)

        buffer = BytesIO(img_bytes)
        img = Image.open(buffer).resize((self.resolution, self.resolution))
        img = self.transform(img)

        return img

• 构造函数 __init__：

  • path：LMDB 数据库的路径。
  • transform：图像转换函数，用于对图像进行预处理。
  • resolution：图像的分辨率，默认为 256x256。
  • max_num：最大加载的图像数量，默认为 70000。

• 初始化过程：

  • 使用 lmdb.open 打开 LMDB 数据库，并设置一些参数如 max_readers、readonly 等。
  • 如果无法打开数据库，抛出 IOError 异常。
  • 通过 self.env.begin(write=False) 获取事务，并遍历数据库的所有键（图像的唯一标识）。
  • 将键存储在 self.keys 列表中，同时限制加载的图像数量为 max_num。
  • 设置 self.length 为数据集的长度（图像数量）、保存 resolution 和 transform 函数。

• __len__ 方法：

  • 返回数据集的长度，即 self.length，表示图像的总数。

• __getitem__ 方法：

  • 根据索引 index 从 LMDB 数据库中读取对应的图像数据。
  • 使用 self.env.begin(write=False) 打开事务，并通过键 key 获取图像的字节数据 img_bytes。
  • 将图像字节数据读取到 BytesIO 缓冲区中，并使用 PIL.Image.open 打开图像。
  • 调整图像大小为 self.resolution x self.resolution，然后应用 transform 进行图像预处理。
  • 返回预处理后的图像 img。

2、NormalDataset 类

功能是从普通的图像文件夹加载图像数据。

class NormalDataset(Dataset):
    def __init__(self, path, transform, resolution=256, max_num=70000):
        self.files = []
        listed_files = sorted(list(list_files(path)))
        for i in range(min(max_num, len(listed_files))):
            file = listed_files[i]
            if any(file.lower().endswith(ext) for ext in IMG_EXTENSIONS):
                self.files.append(file)

        self.resolution = resolution
        self.transform = transform
        self.length = len(self.files)

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        img = Image.open(self.files[index]).resize((self.resolution, self.resolution))
        img = self.transform(img)

        return img

• 构造函数 __init__：

  • path：图像文件夹的路径。
  • transform：图像转换函数，用于对图像进行预处理。
  • resolution：图像的分辨率，默认为 256x256。
  • max_num：最大加载的图像数量，默认为 70000。

• 初始化过程：

  • 使用 imutils.paths.list_files 函数列出指定路径下的所有文件，并按名称排序。
  • 遍历列表中的文件，如果文件的扩展名在 IMG_EXTENSIONS 中（支持的图像格式），则将文件路径添加到 self.files 列表中。
  • 设置 self.resolution、self.transform 和 self.length 分别为图像分辨率、转换函数和数据集的长度。

• __len__ 方法：

  • 返回数据集的长度，即 self.length，表示图像的总数。

• __getitem__ 方法：

  • 根据索引 index 从 self.files 列表中读取对应的图像文件路径。
  • 使用 PIL.Image.open 打开图像，并调整大小为 self.resolution x self.resolution。
  • 应用 transform 进行图像预处理。
  • 返回预处理后的图像 img。

3、set_dataset 函数

功能是根据指定的数据集类型 type 创建相应的数据集实例。

def set_dataset(type, path, transform, resolution):
    datatype = None
    if type == 'lmdb':
        datatype = LMDBDataset
    elif type == 'normal':
        datatype = NormalDataset
    else:
        raise NotImplementedError
    return datatype(path, transform, resolution)

• 参数：

  • type：字符串，指定数据集类型，可以是 'lmdb' 或 'normal'。
  • path：数据集的路径。
  • transform：图像转换函数。
  • resolution：图像分辨率。

• 返回值：

  • 根据 type 的不同值选择并返回对应的数据集实例。

• 异常处理：

  • 如果 type 不是 'lmdb' 或 'normal'，则抛出 NotImplementedError 异常。

（四）utils.py

1、时间格式转换 (time_change)

• 功能：将秒数转换为小时、分钟和秒的字符串格式。
• 细节：
  • 根据输入的 time_init 计算小时、分钟和秒数。
  • 将计算结果以字符串形式拼接成格式化的时间字符串并返回。

def time_change(time_init):
    time_list = []
    if time_init / 3600 > 1:
        time_h = int(time_init / 3600)
        time_m = int((time_init - time_h * 3600) / 60)
        time_s = int(time_init - time_h * 3600 - time_m * 60)
        time_list.append(str(time_h))
        time_list.append('h ')
        time_list.append(str(time_m))
        time_list.append('m ')

    elif time_init / 60 > 1:
        time_m = int(time_init / 60)
        time_s = int(time_init - time_m * 60)
        time_list.append(str(time_m))
        time_list.append('m ')
    else:
        time_s = int(time_init)

    time_list.append(str(time_s))
    time_list.append('s')
    time_str = ''.join(time_list)
    return time_str

2、训练工具 (data_sampler, requires_grad, accumulate, sample_data)

• data_sampler(dataset, shuffle)：

  • 根据 shuffle 参数返回对应的数据采样器，可以是随机采样器或顺序采样器。

• requires_grad(model, flag=True)：

  • 设置模型 model 中所有参数的 requires_grad 属性为 flag。

• accumulate(model1, model2, decay=0.999)：

  • 对 model1 的参数进行滑动平均更新，更新系数为 decay，通过与 model2 的参数进行加权平均实现。

• sample_data(loader)：

  • 从数据加载器 loader 中无限循环地产生数据批次。

'''
Training tools
'''
def data_sampler(dataset, shuffle):
    if shuffle:
        return data.RandomSampler(dataset)

    else:
        return data.SequentialSampler(dataset)


def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag


def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch

3、二进制消息与张量之间的映射 (message_to_tensor, tensor_to_message)

• message_to_tensor(message, sigma, delta)：

  • 将二进制消息 message 转换为张量。
  • sigma 控制消息的编码密度，delta 控制噪声的幅度。
  • 返回一个张量，表示编码后的消息。

• tensor_to_message(secret_tensor, sigma)：

  • 将张量 secret_tensor 反向解码为二进制消息。
  • 返回一个二进制消息的张量。

'''
Mapping between binary messages and floating tensors
'''

def message_to_tensor(message, sigma, delta):
    secret_tensor = torch.zeros(size=(message.shape[0], message.shape[1] // sigma))
    step = 2 / 2 ** sigma
    random_interval_size = step * delta
    message_nums = torch.zeros_like(secret_tensor)
    for i in range(sigma):
        message_nums += message[:, i::sigma] * 2 ** (sigma - i - 1)
    secret_tensor = step * (message_nums + 0.5) - 1
    secret_tensor = secret_tensor + (torch.rand_like(secret_tensor) * random_interval_size * 2 - random_interval_size)
    return secret_tensor


def tensor_to_message(secret_tensor, sigma):
    message = torch.zeros(size=(secret_tensor.shape[0], secret_tensor.shape[1] * sigma))
    step = 2 / 2 ** sigma
    secret_tensor = torch.clamp(secret_tensor, min=-1, max=1) + 1
    message_nums = secret_tensor / step
    zeros = torch.zeros_like(message_nums)
    ones = torch.ones_like(message_nums)
    for i in range(sigma):
        zero_one_map = torch.where(message_nums >= 2 ** (sigma - i - 1), ones, zeros)
        message[:, i::sigma] = zero_one_map
        message_nums -= zero_one_map * 2 ** (sigma - i - 1)
    return message

4、损失函数 (d_logistic_loss, d_r1_loss, g_nonsaturating_loss)

• d_logistic_loss(real_pred, fake_pred)：

  • 计算生成对抗网络的 logistic 损失函数，用于判别器。
  • 返回实部和虚部的损失均值。

• d_r1_loss(real_pred, real_img)：

  • 计算梯度惩罚损失函数，用于判别器。
  • 返回梯度惩罚的值。

• g_nonsaturating_loss(fake_pred)：

  • 计算生成对抗网络的非饱和损失函数，用于生成器。
  • 返回损失值。

'''
Loss functions
'''

def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)

    return real_loss.mean() + fake_loss.mean()


def d_r1_loss(real_pred, real_img):
    (grad_real,) = autograd.grad(
        outputs=real_pred.sum(), inputs=real_img, create_graph=True
    )
    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()

    return grad_penalty


def g_nonsaturating_loss(fake_pred):
    loss = F.softplus(-fake_pred).mean()

    return loss

5、图像裁剪操作 (patchify_image)

patchify_image(img, n_crop, min_size=1/8, max_size=1/4)：

• 将输入的图像 img 裁剪成多个尺寸不同的小块。
• n_crop 控制裁剪的数量，min_size 和 max_size 控制裁剪块的最小和最大尺寸。
• 返回裁剪后的图像小块的张量。

def patchify_image(img, n_crop, min_size=1 / 8, max_size=1 / 4):
    crop_size = torch.rand(n_crop) * (max_size - min_size) + min_size
    batch, channel, height, width = img.shape
    target_h = int(height * max_size)
    target_w = int(width * max_size)
    crop_h = (crop_size * height).type(torch.int64).tolist()
    crop_w = (crop_size * width).type(torch.int64).tolist()

    patches = []
    for c_h, c_w in zip(crop_h, crop_w):
        c_y = random.randrange(0, height - c_h)
        c_x = random.randrange(0, width - c_w)

        cropped = img[:, :, c_y: c_y + c_h, c_x: c_x + c_w]
        cropped = F.interpolate(
            cropped, size=(target_h, target_w), mode="bilinear", align_corners=False
        )

        patches.append(cropped)

    patches = torch.stack(patches, 1).view(-1, channel, target_h, target_w)

    return patches

（五）stylegan2-model.py

这部分和前面的models.py需要区分开（gpt已经生成了，就顺带贴在这里了）

1、PixelNorm 模块

# PixelNorm 模块
class PixelNorm(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input):
        return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)

PixelNorm类是一个自定义的PyTorch模块，用于对输入特征图进行像素归一化。该模块没有可学习参数，仅在前向传播中进行操作。

• __init__ 方法：初始化模块，调用父类的初始化方法。
• forward 方法：实现前向传播。
  • 计算输入张量的平方并在第一个维度上求均值（dim=1）。
  • 加上一个非常小的值1e-8以防止除零错误。
  • 对均值取平方根的倒数（rsqrt），并与输入张量逐元素相乘，实现归一化。

2、make_kernel 函数

def make_kernel(k):
    k = torch.tensor(k, dtype=torch.float32)

    if k.ndim == 1:
        k = k[None, :] * k[:, None]

    k /= k.sum()

    return k

make_kernel函数将一维或二维的核（kernel）转换为二维卷积核，并进行归一化处理。

• 将输入的核转换为torch.float32类型的张量。
• 如果输入是一维的，将其转换为二维的外积形式。
• 对核进行归一化，使其元素和为1。

3. Upsample 模块 

class Upsample(nn.Module):
    def __init__(self, kernel, factor=2):
        super().__init__()

        self.factor = factor
        kernel = make_kernel(kernel) * (factor ** 2)
        self.register_buffer("kernel", kernel)

        p = kernel.shape[0] - factor

        pad0 = (p + 1) // 2 + factor - 1
        pad1 = p // 2

        self.pad = (pad0, pad1)

    def forward(self, input):
        out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)

        return out

Upsample类用于对输入特征图进行上采样操作。

• __init__ 方法：初始化模块。
  • 调用make_kernel函数生成并归一化卷积核，并乘以factor的平方以调整核的强度。
  • 使用register_buffer方法将卷积核注册为不可学习的缓冲区。
  • 计算填充（padding）的大小，并保存为实例变量self.pad。
• forward 方法：实现前向传播。
  • 调用upfirdn2d函数进行上采样，传入输入特征图、卷积核、上采样因子、下采样因子和填充。

4. Downsample 模块

class Downsample(nn.Module):
    def __init__(self, kernel, factor=2):
        super().__init__()

        self.factor = factor
        kernel = make_kernel(kernel)
        self.register_buffer("kernel", kernel)

        p = kernel.shape[0] - factor

        pad0 = (p + 1) // 2
        pad1 = p // 2

        self.pad = (pad0, pad1)

    def forward(self, input):
        out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)

        return out

Downsample类用于对输入特征图进行下采样操作。

• __init__ 方法：初始化模块。
  • 调用make_kernel函数生成并归一化卷积核。
  • 使用register_buffer方法将卷积核注册为不可学习的缓冲区。
  • 计算填充（padding）的大小，并保存为实例变量self.pad。
• forward 方法：实现前向传播。
  • 调用upfirdn2d函数进行下采样，传入输入特征图、卷积核、上采样因子、下采样因子和填充。

5、Blur 模块

class Blur(nn.Module):
    def __init__(self, kernel, pad, upsample_factor=1):
        super().__init__()

        kernel = make_kernel(kernel)

        if upsample_factor > 1:
            kernel = kernel * (upsample_factor ** 2)

        self.register_buffer("kernel", kernel)

        self.pad = pad

    def forward(self, input):
        out = upfirdn2d(input, self.kernel, pad=self.pad)

        return out

• Blur 类继承自 nn.Module，用于应用模糊操作。
• __init__ 方法初始化模块，接受 kernel（核大小）、pad（填充大小）、upsample_factor（上采样因子，默认为 1）作为参数。
• make_kernel 函数用于生成核。
• 如果 upsample_factor 大于 1，则扩展核的大小。
• 使用 register_buffer 方法注册 kernel 为模块的缓冲区。
• self.pad 保存填充大小。
• forward 方法实现模块的前向传播，调用了 upfirdn2d 函数，对输入 input 应用二维上采样卷积操作，并返回结果 out。

6、EqualConv2d 模块

class EqualConv2d(nn.Module):
    def __init__(
        self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
    ):
        super().__init__()

        self.weight = nn.Parameter(
            torch.randn(out_channel, in_channel, kernel_size, kernel_size)
        )
        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)

        self.stride = stride
        self.padding = padding

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_channel))

        else:
            self.bias = None

    def forward(self, input):
        out = F.conv2d(
            input,
            self.weight * self.scale,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
        )

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
            f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
        )

EqualConv2d 类继承自 nn.Module，实现了具有权重均匀初始化的二维卷积操作。 

• __init__ 方法初始化模块，接受 in_channel（输入通道数）、out_channel（输出通道数）、kernel_size（核大小）、stride（步长，默认为 1）、padding（填充大小，默认为 0）、bias（是否使用偏置，默认为 True）作为参数。
• 使用 nn.Parameter 创建权重 self.weight，其形状为 (out_channel, in_channel, kernel_size, kernel_size)，使用标准正态分布随机初始化。
• 计算 self.scale 作为权重的缩放系数，用于保证权重初始化的数值稳定性。
• 设置 self.stride 和 self.padding 为类属性。
• 如果 bias 为 True，则创建偏置 self.bias，其形状为 (out_channel,)，初始化为零。
• forward 方法实现模块的前向传播，调用了 PyTorch 中的 F.conv2d 函数，对输入 input 应用二维卷积操作，并返回结果 out。
• __repr__ 方法定义了模块的字符串表示形式，用于打印模块的信息，包括权重的形状、步长和填充大小。

7、EqualLinear 模块

class EqualLinear(nn.Module):
    def __init__(
        self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None
    ):
        super().__init__()

        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))

        else:
            self.bias = None

        self.activation = activation

        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
        self.lr_mul = lr_mul

    def forward(self, input):
        if self.activation:
            out = F.linear(input, self.weight * self.scale)
            out = fused_leaky_relu(out, self.bias * self.lr_mul)

        else:
            out = F.linear(
                input, self.weight * self.scale, bias=self.bias * self.lr_mul
            )

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"
        )

EqualLinear 类继承自 nn.Module，实现了具有均匀初始化权重的全连接层。

• __init__ 方法初始化模块，接受 in_dim（输入维度）、out_dim（输出维度）、bias（是否使用偏置，默认为 True）、bias_init（偏置初始值，默认为 0）、lr_mul（学习率倍增因子，默认为 1）、activation（激活函数，默认为 None）作为参数。
• 使用 nn.Parameter 创建权重 self.weight，其形状为 (out_dim, in_dim)，并进行缩放。
• 如果 bias 为 True，则创建偏置 self.bias，其形状为 (out_dim,)，初始化为 bias_init。
• self.activation 保存激活函数。
• 计算 self.scale 作为权重的缩放系数，用于保证权重初始化的数值稳定性。
• 设置 self.lr_mul 为类属性。
• forward 方法实现模块的前向传播，如果存在激活函数，则先进行线性变换，再应用激活函数；否则直接进行线性变换。
• __repr__ 方法定义了模块的字符串表示形式，用于打印模块的信息，包括权重的形状。

8、ScaledLeakyReLU 模块

class ScaledLeakyReLU(nn.Module):
    def __init__(self, negative_slope=0.2):
        super().__init__()

        self.negative_slope = negative_slope

    def forward(self, input):
        out = F.leaky_relu(input, negative_slope=self.negative_slope)

        return out * math.sqrt(2)

ScaledLeakyReLU 类继承自 nn.Module，实现了带有缩放系数的 Leaky ReLU 激活函数。

• __init__ 方法初始化模块，接受 negative_slope（负斜率，默认为 0.2）作为参数。
• self.negative_slope 保存负斜率。
• forward 方法实现模块的前向传播，调用了 PyTorch 中的 F.leaky_relu 函数，对输入 input 应用 Leaky ReLU 激活函数，并进行缩放。
• 返回值为应用激活函数和缩放后的输出。

9、ModulatedConv2d 模块

class ModulatedConv2d(nn.Module):
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        style_dim,
        demodulate=True,
        upsample=False,
        downsample=False,
        blur_kernel=[1, 3, 3, 1],
    ):
        super().__init__()

        self.eps = 1e-8
        self.kernel_size = kernel_size
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.upsample = upsample
        self.downsample = downsample

        if upsample:
            factor = 2
            p = (len(blur_kernel) - factor) - (kernel_size - 1)
            pad0 = (p + 1) // 2 + factor - 1
            pad1 = p // 2 + 1

            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)

        if downsample:
            factor = 2
            p = (len(blur_kernel) - factor) + (kernel_size - 1)
            pad0 = (p + 1) // 2
            pad1 = p // 2

            self.blur = Blur(blur_kernel, pad=(pad0, pad1))

        fan_in = in_channel * kernel_size ** 2
        self.scale = 1 / math.sqrt(fan_in)
        self.padding = kernel_size // 2

        self.weight = nn.Parameter(
            torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)
        )

        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)

        self.demodulate = demodulate

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, "
            f"upsample={self.upsample}, downsample={self.downsample})"
        )

    def forward(self, input, style):
        batch, in_channel, height, width = input.shape

        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
        weight = self.scale * self.weight * style

        if self.demodulate:
            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)

        weight = weight.view(
            batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size
        )

        if self.upsample:
            input = input.view(1, batch * in_channel, height, width)
            weight = weight.view(
                batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size
            )
            weight = weight.transpose(1, 2).reshape(
                batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
            )
            out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch)
            _, _, height, width = out.shape
            out = out.view(batch, self.out_channel, height, width)
            out = self.blur(out)

        elif self.downsample:
            input = self.blur(input)
            _, _, height, width = input.shape
            input = input.view(1, batch * in_channel, height, width)
            out = F.conv2d(input, weight, padding=0, stride=2, groups=batch)
            _, _, height, width = out.shape
            out = out.view(batch, self.out_channel, height, width)

        else:
            input = input.view(1, batch * in_channel, height, width)
            out = F.conv2d(input, weight, padding=self.padding, groups=batch)
            _, _, height, width = out.shape
            out = out.view(batch, self.out_channel, height, width)

        return out

ModulatedConv2d 类继承自 nn.Module，实现了具有调制功能的二维卷积层。

• __init__ 方法初始化模块，接受多个参数：in_channel（输入通道数）、out_channel（输出通道数）、kernel_size（核大小）、style_dim（样式维度）、demodulate（是否去调制，默认为 True）、upsample（是否上采样，默认为 False）、downsample（是否下采样，默认为 False）、blur_kernel（模糊核，默认为 [1, 3, 3, 1]）。
• 根据是否上采样或下采样，计算填充大小，并创建模糊层 self.blur。
• 计算 fan_in 和 self.scale 用于权重初始化的缩放。
• 使用 nn.Parameter 创建权重 self.weight，其形状为 (1, out_channel, in_channel, kernel_size, kernel_size)。
• 创建调制层 self.modulation，使用 EqualLinear 实现。
• self.demodulate 保存是否去调制。
• __repr__ 方法定义了模块的字符串表示形式，用于打印模块的信息，包括输入通道数、输出通道数和核大小。
• forward 方法实现模块的前向传播：
  • 计算样式 style 的调制值，并调整权重。
  • 如果需要去调制，则计算去调制因子，并调整权重。
  • 根据是否上采样、下采样或正常卷积，分别进行相应的卷积操作。
  • 返回值为应用卷积和调制后的输出。

10、NoiseInjection 模块

class NoiseInjection(nn.Module):
    def __init__(self):
        super().__init__()

        self.weight = nn.Parameter(torch.zeros(1))

    def forward(self, image, noise=None):
        if noise is None:
            batch, _, height, width = image.shape
            noise = image.new_empty(batch, 1, height, width).normal_()

        return image + self.weight * noise

NoiseInjection 类继承自 nn.Module，实现噪声注入操作。

• NoiseInjection 类继承自 nn.Module，实现了噪声注入操作。
• __init__ 方法初始化模块，创建一个可学习的参数 self.weight，并初始化为零。
• forward 方法实现模块的前向传播：
  • 接受输入图像 image 和可选的噪声 noise 作为参数。
  • 如果没有提供噪声 noise，则根据输入图像的形状生成一个与输入图像大小相同的随机噪声。
  • 将输入图像和噪声按 self.weight 进行加权和，并返回结果。

11、ConstantInput 模块

class ConstantInput(nn.Module):
    def __init__(self, channel, size=4):
        super().__init__()

        self.input = nn.Parameter(torch.randn(1, channel, size, size))

    def forward(self, input):
        batch = input.shape[0]
        out = self.input.repeat(batch, 1, 1, 1)

        return out

ConstantInput 类继承自 nn.Module，实现了一个常数输入的模块。

• __init__ 方法初始化模块，接受 channel（通道数）和 size（大小，默认为 4）作为参数。
• 创建一个形状为 (1, channel, size, size) 的可学习参数 self.input，并初始化为标准正态分布的随机值。
• forward 方法实现模块的前向传播：
  • 返回结果。
  • 将 self.input 重复批量大小次，以匹配输入的批量大小。
  • 接受输入 input 作为参数，用于获取批量大小。

12、StyledConv 模块

class StyledConv(nn.Module):
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        style_dim,
        upsample=False,
        blur_kernel=[1, 3, 3, 1],
        demodulate=True,
    ):
        super().__init__()

        self.conv = ModulatedConv2d(
            in_channel,
            out_channel,
            kernel_size,
            style_dim,
            upsample=upsample,
            blur_kernel=blur_kernel,
            demodulate=demodulate,
        )

        self.noise = NoiseInjection()
        # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
        # self.activate = ScaledLeakyReLU(0.2)
        self.activate = FusedLeakyReLU(out_channel)

    def forward(self, input, style, noise=None):
        out = self.conv(input, style)
        out = self.noise(out, noise=noise)
        # out = out + self.bias
        out = self.activate(out)

        return out

StyledConv 类继承自 nn.Module，实现了一个带有样式调制的卷积层。

• __init__ 方法初始化模块，接受多个参数：in_channel（输入通道数）、out_channel（输出通道数）、kernel_size（核大小）、style_dim（样式维度）、upsample（是否上采样，默认为 False）、blur_kernel（模糊核，默认为 [1, 3, 3, 1]）、demodulate（是否去调制，默认为 True）。
• 创建一个 ModulatedConv2d 模块 self.conv，用于样式调制的卷积操作。
• 创建一个 NoiseInjection 模块 self.noise，用于注入噪声。
• 创建一个 FusedLeakyReLU 模块 self.activate，用于激活操作（注意，这里用的是 FusedLeakyReLU，注释掉了 ScaledLeakyReLU 和 bias）。
• forward 方法实现模块的前向传播：
  • 接受输入 input、样式 style 和可选的噪声 noise 作为参数。
  • 调用 self.conv 对输入进行样式调制的卷积操作。
  • 调用 self.noise 对卷积输出添加噪声。
  • 调用 self.activate 对添加噪声后的输出进行激活操作。
  • 返回结果。

13、StyledConv_without_noise 模块

StyledConv_without_noise 类似于之前定义的 StyledConv 模块，但不包含噪声注入部分。

class StyledConv_without_noise(nn.Module):
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        style_dim,
        upsample=False,
        blur_kernel=[1, 3, 3, 1],
        demodulate=True,
    ):
        super().__init__()

        self.conv = ModulatedConv2d(
            in_channel,
            out_channel,
            kernel_size,
            style_dim,
            upsample=upsample,
            blur_kernel=blur_kernel,
            demodulate=demodulate,
        )

        #self.noise = NoiseInjection()
        # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
        # self.activate = ScaledLeakyReLU(0.2)
        self.activate = FusedLeakyReLU(out_channel)

    def forward(self, input, style, noise=None):
        out = self.conv(input, style)
        #out = self.noise(out, noise=noise)
        # out = out + self.bias
        out = self.activate(out)

        return out

• __init__ 方法:

  • 接受多个参数：in_channel（输入通道数）、out_channel（输出通道数）、kernel_size（卷积核大小）、style_dim（样式维度）、upsample（是否上采样，默认为 False）、blur_kernel（模糊核，默认为 [1, 3, 3, 1]）、demodulate（是否去调制，默认为 True）。
  • 创建一个 ModulatedConv2d 实例 self.conv，用于样式调制的卷积操作。
  • 创建一个 FusedLeakyReLU 实例 self.activate，用于激活操作。

• forward 方法:

  • 接受输入 input 和样式 style 作为参数。
  • 使用 self.conv 对输入进行样式调制的卷积操作。
  • 使用 self.activate 对卷积输出进行激活操作。
  • 返回激活后的结果。

14、ToRGB 模块

class ToRGB(nn.Module):
    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
        super().__init__()

        if upsample:
            self.upsample = Upsample(blur_kernel)

        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))

    def forward(self, input, style, skip=None):
        out = self.conv(input, style)
        out = out + self.bias

        if skip is not None:
            skip = self.upsample(skip)

            out = out + skip

        return out

ToRGB 模块用于将特征图转换为 RGB 图像。

• __init__ 方法:

  • 接受多个参数：in_channel（输入通道数）、style_dim（样式维度）、upsample（是否上采样，默认为 True）、blur_kernel（模糊核，默认为 [1, 3, 3, 1]）。
  • 如果 upsample 为 True，则创建一个 Upsample 实例 self.upsample。
  • 创建一个 ModulatedConv2d 实例 self.conv，用于样式调制的卷积操作，输出通道数为 3。
  • 创建一个可学习的偏置参数 self.bias，并初始化为零。

• forward 方法:

  • 接受输入 input、样式 style 和可选的 skip 作为参数。
  • 使用 self.conv 对输入进行样式调制的卷积操作。
  • 将卷积输出与 self.bias 相加。
  • 如果提供了 skip，则对 skip 进行上采样，并与卷积输出相加。
  • 返回结果。

15、Generator 模块

Generator 模块实现了整个生成器网络。

class Generator(nn.Module):
    # __init__方法：接受多个参数：size（生成图像的大小）、style_dim（样式维度）、n_mlp（MLP 层数）、channel_multiplier（通道数乘数，默认为 2）、blur_kernel（模糊核，默认为 [1, 3, 3, 1]）、lr_mlp（MLP 的学习率乘数，默认为 0.01）。
    def __init__(
        self,
        size,
        style_dim,
        n_mlp,
        channel_multiplier=2,
        blur_kernel=[1, 3, 3, 1],
        lr_mlp=0.01,
    ):
        super().__init__()

        self.size = size
        
        self.style_dim = style_dim
        # # 创建一个 PixelNorm 实例，并添加到 layers 列表中。
        layers = [PixelNorm()]
        # 使用 EqualLinear 创建多个 MLP 层，并添加到 layers 列表中，最后将 layers 列表构造成一个 nn.Sequential 实例 self.style。
        for i in range(n_mlp):
            layers.append(
                EqualLinear(
                    style_dim, style_dim, lr_mul=lr_mlp, activation="fused_lrelu"
                )
            )
        self.style = nn.Sequential(*layers)

        # 定义不同分辨率下的通道数，并存储在 self.channels 字典中。
        self.channels = {
            4: 512,
            8: 512,
            16: 512,
            32: 512,
            64: 256 * channel_multiplier,
            128: 128 * channel_multiplier,
            256: 64 * channel_multiplier,
            512: 32 * channel_multiplier,
            1024: 16 * channel_multiplier,
        }
        # 创建一个 ConstantInput 实例 self.input。
        self.input = ConstantInput(self.channels[4])

        # 创建第一个卷积层 self.conv1 和第一个 ToRGB 层 self.to_rgb1。
        self.conv1 = StyledConv(
            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
        )
        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)

        # 计算日志大小和层数，并存储在 self.log_size 和 self.num_layers 中。
        self.log_size = int(math.log(size, 2))
        self.num_layers = (self.log_size - 2) * 2 + 1

        # 创建卷积层列表 self.convs、上采样层列表 self.upsamples、ToRGB 层列表 self.to_rgbs 和噪声层列表 self.noises。
        self.convs = nn.ModuleList()
        self.upsamples = nn.ModuleList()
        self.to_rgbs = nn.ModuleList()
        self.noises = nn.Module()

        in_channel = self.channels[4]

        for layer_idx in range(self.num_layers):
            res = (layer_idx + 5) // 2
            shape = [1, 1, 2 ** res, 2 ** res]
            self.noises.register_buffer(f"noise_{layer_idx}", torch.randn(*shape))

        for i in range(3, self.log_size + 1):
            out_channel = self.channels[2 ** i]

            self.convs.append(
                StyledConv(
                    in_channel,
                    out_channel,
                    3,
                    style_dim,
                    upsample=True,
                    blur_kernel=blur_kernel,
                )
            )

            self.convs.append(
                StyledConv(
                    out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
                )
            )

            self.to_rgbs.append(ToRGB(out_channel, style_dim))

            in_channel = out_channel

        self.n_latent = self.log_size * 2 - 2

    # 生成一组随机噪声，用于不同分辨率下的噪声注入
    def make_noise(self):
        device = self.input.input.device

        noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]

        for i in range(3, self.log_size + 1):
            for _ in range(2):
                noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))

        return noises

    # 计算多个潜在向量的平均值，用于样式平均
    def mean_latent(self, n_latent):
        latent_in = torch.randn(
            n_latent, self.style_dim, device=self.input.input.device
        )
        latent = self.style(latent_in).mean(0, keepdim=True)

        return latent

    # 将输入映射到样式空间
    def get_latent(self, input):
        return self.style(input)

    # 实现生成器网络的前向传播，接受多个参数：
    # styles：样式向量。                   return_latents：是否返回潜在向量。
    # inject_index：样式混合的插入索引。   truncation：截断超参数。
    # truncation_latent：截断潜在向量。    input_is_latent：输入是否为潜在向量。
    # noise：噪声。                       randomize_noise：是否随机化噪声。
    def forward(
        self,
        styles,
        return_latents=False,
        inject_index=None,
        truncation=1,
        truncation_latent=None,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
    ):
        # 根据 input_is_latent 判断是否需要将输入映射到样式空间
        if not input_is_latent:
            styles = [self.style(s) for s in styles]

        if noise is None:
            if randomize_noise:
                noise = [None] * self.num_layers
            else:
                noise = [
                    getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
                ]
        # 根据 truncation 和 truncation_latent 进行样式截断
        if truncation < 1:
            style_t = []

            for style in styles:
                style_t.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )

            styles = style_t
        # 处理样式混合逻辑
        if len(styles) < 2:
            inject_index = self.n_latent

            if styles[0].ndim < 3:
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)

            else:
                latent = styles[0]

        else:
            if inject_index is None:
                inject_index = random.randint(1, self.n_latent - 1)

            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
            latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)

            latent = torch.cat([latent, latent2], 1)

        # 按层级结构进行卷积操作、上采样操作和 ToRGB 操作。

        # 将潜在向量 latent 输入到 ConstantInput 层，生成初始输入特征图。
        out = self.input(latent)
        # 使用第一个 StyledConv 层对输入特征图进行卷积操作，并应用样式latent[:, 0] 和噪声 noise[0]。
        out = self.conv1(out, latent[:, 0], noise=noise[0])
        # 将初始卷积输出 out 转换为 RGB 图像，并应用样式 latent[:, 1]。结果存储在 skip 中。
        skip = self.to_rgb1(out, latent[:, 1])

        i = 1
        # 使用 for 循环遍历 self.convs（卷积层列表）、noise（噪声列表）和 self.to_rgbs（ToRGB 层列表）。
        for conv1, conv2, noise1, noise2, to_rgb in zip(
        # self.convs 按照两个一组的方式进行处理，因此使用 self.convs[::2] 和 self.convs[1::2] 来获取偶数和奇数索引的卷积层
            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
        ):
        # conv1 卷积操作：应用样式 latent[:, i] 和噪声 noise1；conv2 卷积操作：应用样式 latent[:, i + 1] 和噪声 noise2；to_rgb 操作：将卷积输出转换为 RGB 图像，并应用样式 latent[:, i + 2] 和 skip 图像。
            out = conv1(out, latent[:, i], noise=noise1)
            out = conv2(out, latent[:, i + 1], noise=noise2)
            skip = to_rgb(out, latent[:, i + 2], skip)

            i += 2
        # 循环结束后，最终生成的图像存储在 image 中。
        image = skip
        # 返回生成的图像和潜在向量（如果需要）。
        if return_latents:
            return image, latent

        else:
            return image, None

16、ConvLayer模块

该模块继承自 nn.Sequential。这个类的主要作用是根据给定的参数构建一个包含卷积操作、可选的下采样（downsample）操作、以及激活函数（activate）的卷积层。

class ConvLayer(nn.Sequential):
    # in_channel：输入通道数。      out_channel：输出通道数。
    # kernel_size：卷积核大小。     downsample：是否进行下采样，默认值为 False。
    # blur_kernel：用于模糊操作的内核，默认值为 [1, 3, 3, 1]。
    # bias：是否在卷积操作中使用偏置，默认值为 True。
    # activate：是否添加激活函数，默认值为 True。
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        downsample=False,
        blur_kernel=[1, 3, 3, 1],
        bias=True,
        activate=True,
    ):
        # 定义卷积层列表
        layers = []

        # 如果是下采样操作
        if downsample:
        # 计算 pad0 和 pad1 以确定填充
            factor = 2
            p = (len(blur_kernel) - factor) + (kernel_size - 1)
            pad0 = (p + 1) // 2
            pad1 = p // 2
        # 添加模糊层 Blur
            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
        # 设置卷积步幅 stride 为 2，并将填充设置为 0
            stride = 2
            self.padding = 0

        else:
        # 如果不是下采样操作
        # 设置卷积步幅 stride 为 1，并计算填充
            stride = 1
            self.padding = kernel_size // 2

        # 将卷积层 EqualConv2d 添加到 layers 列表中
        layers.append(
            EqualConv2d(
                in_channel,
                out_channel,
                kernel_size,
                padding=self.padding,
                stride=stride,
                bias=bias and not activate,
            )
        )
        # 如果 activate 为 True，则添加激活函数到 layers 列表中
        # FusedLeakyReLU：如果 bias 为 True，则使用融合的 Leaky ReLU 激活函数。
        # ScaledLeakyReLU：如果 bias 为 False，则使用缩放的 Leaky ReLU 激活函数。
        if activate:
            if bias:
                layers.append(FusedLeakyReLU(out_channel))

            else:
                layers.append(ScaledLeakyReLU(0.2))
        # 调用父类 nn.Sequential 的构造函数，并将 layers 列表传递给它
        super().__init__(*layers)

17、ResBlock模块

ResBlock 继承自 nn.Module，表示一个残差块，它包含两个卷积层和一个跳跃连接（skip connection）。

class ResBlock(nn.Module):
    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
        super().__init__()

        self.conv1 = ConvLayer(in_channel, in_channel, 3)
        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)

        self.skip = ConvLayer(
            in_channel, out_channel, 1, downsample=True, activate=False, bias=False
        )

    def forward(self, input):
        out = self.conv1(input)
        out = self.conv2(out)

        skip = self.skip(input)
        out = (out + skip) / math.sqrt(2)

        return out

初始化方法 __init__

• 参数：

  • in_channel：输入通道数。
  • out_channel：输出通道数。
  • blur_kernel：用于模糊操作的内核，默认值为 [1, 3, 3, 1]。

• 卷积层 conv1 和 conv2：

  • conv1：一个普通的卷积层。
  • conv2：一个带有下采样操作的卷积层。

• 跳跃连接 skip：

  • 一个 1x1 卷积层，也带有下采样操作，但没有激活函数和偏置。

前向传播方法 forward

• 输入通过 conv1 和 conv2 进行卷积操作。
• 跳跃连接通过 skip 进行卷积操作。
• 最后，将 conv2 的输出和跳跃连接的输出相加，并除以 sqrt{2}。

18、Discriminator模块

Discriminator 继承自 nn.Module，表示一个鉴别器。它包含多个残差块和一些卷积层，用于判别输入图像。

class Discriminator(nn.Module):
    # size：输入图像的尺寸（假设为正方形）。channel_multiplier：通道数的倍增因子。
    # blur_kernel：用于模糊操作的内核，默认值为 [1, 3, 3, 1]。
    def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
        super().__init__()
    # 通道数字典 channels：为不同尺寸的图像定义了不同的通道数。
        channels = {
            4: 512,
            8: 512,
            16: 512,
            32: 512,
            64: 256 * channel_multiplier,
            128: 128 * channel_multiplier,
            256: 64 * channel_multiplier,
            512: 32 * channel_multiplier,
            1024: 16 * channel_multiplier,
        }
        # 卷积层列表 convs：第一个卷积层将输入的 3 通道图像转换为相应的通道数。
        convs = [ConvLayer(3, channels[size], 1)]
        # 根据输入图像的尺寸添加多个残差块 ResBlock 和卷积层
        log_size = int(math.log(size, 2)) #计算输入图像的对数尺寸
        # 初始化输入通道数
        in_channel = channels[size]
        # 根据输入图像的尺寸添加残差块 ResBlock 和卷积层
        for i in range(log_size, 2, -1):
            out_channel = channels[2 ** (i - 1)]

            convs.append(ResBlock(in_channel, out_channel, blur_kernel))

            in_channel = out_channel
        # 将所有卷积层和残差块组合成一个顺序容器 nn.Sequential，以便在前向传播中依次执行。
        self.convs = nn.Sequential(*convs)

        # stddev_group 和 stddev_feat 用于计算标准差。
        self.stddev_group = 4
        self.stddev_feat = 1

        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
        self.final_linear = nn.Sequential(
            EqualLinear(channels[4] * 4 * 4, channels[4], activation="fused_lrelu"),
            EqualLinear(channels[4], 1),
        )

    def forward(self, input):
        out = self.convs(input)

        batch, channel, height, width = out.shape
       
        group = min(batch, self.stddev_group)
        stddev = out.view(
            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
        )
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(group, 1, height, width)
        out = torch.cat([out, stddev], 1)

        # final_conv：一个卷积层，用于合并标准差特征。
        out = self.final_conv(out)
        out = out.view(batch, -1)
        #final_linear：两个全连接层，用于最终的判别。
        out = self.final_linear(out)

        return out

（六）参数配置train_config.txt

• 实验名称：./experiments/1
• 数据集路径：../classroom_train
• 数据集类型：普通数据集（normal）
• 迭代次数：80000
• TensorExtractor 模型的输出张量中的通道数N：1
• lambda_Ex：10（lambda_Ex 是一个超参数，通常用于控制损失函数中的正则化项的权重。）
• 检查点路径：None
• 学习率：0.002
• 批量大小：1
• 图像尺寸：256
• 真实图像梯度惩罚因子（real_r1）：10
• 纹理图像梯度惩罚因子（texture_r1）：1
• 分布图像梯度惩罚因子（dist_r1）：1
• 参考图像裁剪数目（ref_crop）：4
• 图像裁剪数目（n_crop）：8
• 每隔多少次迭代进行判别器正则化（d_reg_every）：16
• 通道数（channel）：32
• 通道倍增因子（channel_multiplier）：1
• 结构通道数（structure_channel）：8
• 纹理通道数（texture_channel）：2048
• 每隔多少次迭代记录日志（log_every）：200
• 每隔多少次迭代显示结果（show_every）：1000
• 每隔多少次迭代保存模型（save_every）：200000
• 起始迭代次数：0
• 模糊核大小：(1, 3, 3, 1)