Fork me on GitHub

训练自己的户型图扩散模型（附远程服务器配置）

字数统计: 3k字   |   阅读时长≈ 12分

前言：上次做了户型图数据集调研，并构建了自己的户型图数据集，这次使用构建好的数据集训练自己的Diffusion Model

本次我训练了四个版本的Diffusion Model for Floor Plan datasets，分别是：

cubicasa5k-Simple-DM

cubicasa5k-Large-DM

RPLAN-DM

Myself-DM（自己从网上爬取的数据集）

训练伊始的踩坑

ok，我的训练从入门快进到了入土，看一下提示出来的信息吧：

1
2
3
4
5
6
7
8
9
10
11
2023-11-29 18:10:57 LAPTOP-56LC5HQ1 __main__[684] INFO [None]: Input params: Namespace(act='gelu', batch_size=2, cfg_scale=3, conditional=False, dataset_path='../datasets/RPLAN', distributed=False, epochs=300, fp16=False, image_size=64, load_model_dir='', lr=0.0003, lr_func='linear', main_gpu=0, num_classes=1, num_vis=-1, num_workers=0, optim='adamw', result_path='../results', resume=False, run_name='RPLAN_model', sample='ddim', save_model_interval=True, seed=0, start_epoch=-1, start_model_interval=-1, vis=True, world_size=2)
2023-11-29 18:10:57 LAPTOP-56LC5HQ1 utils.initializer[684] INFO The seed is initialized, and the seed ID is 0.
2023-11-29 18:10:57 LAPTOP-56LC5HQ1 utils.initializer[684] INFO Init program, it is checking the basic setting.
2023-11-29 18:10:58 LAPTOP-56LC5HQ1 utils.initializer[684] INFO {'is_init': False, 'device_count': 1, 'device_name': 'NVIDIA GeForce RTX 4070 Laptop GPU', 'device_cap': (8, 9), 'device_prop': _CudaDeviceProperties(name='NVIDIA GeForce RTX 4070 Laptop GPU', major=8, minor=9, total_memory=8187MB, multi_processor_count=36)}
2023-11-29 18:10:58 LAPTOP-56LC5HQ1 __main__[684] INFO [cuda]: Successfully Use normal training.
2023-11-29 18:10:59 LAPTOP-56LC5HQ1 __main__[684] INFO [cuda]: Fp32 training.
2023-11-29 18:10:59 LAPTOP-56LC5HQ1 __main__[684] INFO [cuda]: Start training.
2023-11-29 18:10:59 LAPTOP-56LC5HQ1 __main__[684] INFO [cuda]: Start epoch 0:
2023-11-29 18:10:59 LAPTOP-56LC5HQ1 __main__[684] INFO [cuda]: This epoch learning rate is 0.0003
  0%|          | 0/40394 [00:00<?, ?it/s]/home/ubuntu22/anaconda3/envs/cu121py38/lib/python3.8/site-packages/torch/nn/modules/conv.py:456: UserWarning: Applied workaround for CuDNN issue, install nvrtc.so (Triggered internally at ../aten/src/ATen/native/cudnn/Conv_v8.cpp:80.)
  return F.conv2d(input, weight, bias, self.stride,

image-20231129183427474

我手里就一张4070，只有可怜的8GB显存，batch_size调大点，显存直接爆炸，batch_size调小点，训练一轮竟然要40分钟！！而且我之前也没有训练diffusion model的经验，也不知道大概训练多少轮会收敛，以300轮为例，要8天才能训练完，难顶。

主要是也没有做实验室的工作，用实验室的服务器说不过去，现在只能找一些其他的算力资源了，之前白嫖colab，但是这个靠运气，而且白嫖到的卡也一般，和我这4070半斤八两，那就只能租服务器了，正好也记录一下远程服务器的相关配置啥的。

租带卡服务器用于自己的训练

我选择的租服务器平台是： 服务器平台

当前要花钱的，其实也还好，我比较推荐3090，显存大，单精双精都好，性价比挺高的。

image-20231129195333163

我这里打算租一台两张3090的服务器

image-20231129195404694

3块多一个小时，勉强可以接受~

image-20231129195847557

镜像建议用基础镜像，每台机器对cuda和pytorch的兼容程度不一样，注意选择。

image-20231129200028568

ok，这样我们就创建好了

用pycharm连接远程服务器

连接SSH

根据上面给出的SSH登陆填入pycharm中：

image-20231129200441933

配置python解释器

image-20231129200557893

服务器上后很多python环境，我们要选择我们想要的，现在的服务器一般都是miniconda

现在服务器终端中看一下python在哪

image-20231129201015576

然后选取/root/miniconda3/bin/python复制到pycharm中：

image-20231129201050290

同步文件夹

这里要把当前本地的工程文件上传到服务器中，最后自己命名一个文件夹

image-20231129201236276

之后我们就等本地文件上传至服务器中（这个时间还是挺长的）

image-20231129213927114

ok，经过漫长的等待，在服务器里已经有这个项目了：

image-20231129232028529

很少能打这么富裕的仗，记录一下吧~

image-20231130134455663

image-20231130142801680

这个提示是显存不够（看来batch_zise=16对于2张3090还是太暴力了，那么最直接的办法是batch_size调小一点）

image-20231130163215468

谢天谢地，终于训起来了，该说不说，这富裕仗打的就是爽，一个epoch就十几秒！以后高低得整个10张4090的服务器吧~

数据集预处理

上次已经找好数据集了，但是毕竟是在别的任务上的数据集，一般又以下几个问题：

  • 分辨率过大（RPLAN数据集是256*256的，cubicasa5k是完全不规则的）
  • 数据格式不统一
  • 标注的问题

标注问题上次已经解决过了，现在我们手头的是纯图片数据集，这也正是image-to-image的Diffusion Model所需要的，那么下买面就来解决分辨率的问题，这个问题一开始我没有意识到，所以在我最开始的训练是即使把batch_size调到最小、选用fp16后，显存依旧会爆炸，在和老师沟通后了解到Diffusion Model一般是用64*64的，过大的分辨率会先下采样，那么下面就分数据集来说明对其的预处理。

RPLAN数据集

这个数据集本身就是纯图数据集，主要是本身为256*256的，那么我们只需要resize一下就行

1
2
3
4
5
6
7
8
import os
import cv2
import numpy as np
for filename in os.listdir(r"../RPLAN/class_1/"):
    print(filename)
    img = cv2.imdecode(np.fromfile('../RPLAN/class_1/'+filename, dtype=np.uint8), 1)
    img_new = cv2.resize(img,(64,64))
    cv2.imencode('.jpg', img_new)[1].tofile('../RPLAN_resized/class_1/'+filename)

这样就可以直接用了

cubicasa5k数据集

这个数据集比较麻烦，它的文件结构是这样的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
├── cubicasa5k
│       ├── colorful
│   		└── 30
│       		├── F1_original.png
│       		├── F1_scaled.png
│       		└── model.svg
│   		└── ...
│ 
│       ├── high_quality
│   		└── 17
│       		├── F1_original.png
│       		├── F1_scaled.png
│       		└── model.svg
│   		└── ...
│ 
│       ├── high_quality_architectural
│   		└── 1
│       		├── F1_original.png
│       		├── F1_scaled.png
│       		└── model.svg
│   		└── ...

我大概看了一下，其中的png格式的数据质量不一，colorful里的是彩色的，high_quality是黑白的，high_quality_architectural是较为复杂的，而且其中有的是标准的户型图，有的是手稿，所以我决定提取其中model.svg数据，这个矢量化数据还是比较标准的。

提取model.svg数据

把每个文件夹下的model.svg数据提取出来，放在一个文件夹中，然后转成png格式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import os
import shutil

# 输入文件夹路径
# input_folder = './high_quality'
input_folder = './high_quality_architectural'
# input_folder = './colorful'
# 输出文件夹路径
# output_folder = './high_quality_clean'
output_folder = './high_quality_architectural_clean'
# output_folder = './colorful_clean'

# 确保输出文件夹存在，如果不存在则创建
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历输入文件夹中的所有子文件夹
for subdir, dirs, files in os.walk(input_folder):
    for dir_name in dirs:
        dir_path = os.path.join(subdir, dir_name)
        # 找到每个子文件夹中的model.svg文件
        svg_file = os.path.join(dir_path, 'model.svg')
        # 如果找到了model.svg文件
        if os.path.exists(svg_file):
            # 构造输出文件的路径和新文件名
            output_file = os.path.join(output_folder, f'{len(os.listdir(output_folder)) + 1}.svg')
            # 复制model.svg到输出文件夹，并按照新的文件名进行重命名
            shutil.copy(svg_file, output_file)
            print(f'Copied {svg_file} to {output_file}')

转成png格式：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import cairosvg
import os

# inputFolder = "./high_quality_clean"  # 输入的文件夹，里面有svg
inputFolder = "./high_quality_architectural_clean"  # 输入的文件夹，里面有svg
# inputFolder = "./colorful_clean"  # 输入的文件夹，里面有svg

# outputFolder = "./high_quality_clean_png"  # 输出的文件夹，将把结果放到此文件夹中
outputFolder = "./high_quality_architectural_clean_png"  # 输出的文件夹，将把结果放到此文件夹中
# outputFolder = "./colorful_clean_png"  # 输出的文件夹，将把结果放到此文件夹中

# 确保输出文件夹存在，如果不存在则创建
if not os.path.exists(outputFolder):
    os.makedirs(outputFolder)

for root, dirs, files in os.walk(inputFolder):  # 遍历所有的文件
    for f in files:
        svgFile = os.path.join(root, f)  # svg文件名
        if f[-3:] == "svg":  # 确保是svg
            pngFile = outputFolder + "/" + f.replace("svg", "png")  # png文件名
            try:
                cairosvg.svg2png(url=svgFile, write_to=pngFile, dpi=1900, output_width=500)
            except:
                print('error =>' + pngFile)
            finally:
                print('file => ' + pngFile)

再把所有图片都resize成64*64的大小：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from PIL import Image
import os

# 输入文件夹路径
# input_folder = './high_quality_clean_png'
input_folder = './high_quality_architectural_clean_png'
# input_folder = './colorful_clean_png'

# 输出文件夹路径
# output_folder = './high_quality_clean_png_resize'
output_folder = './high_quality_architectural_clean_png_resize'
# output_folder = './colorful_clean_png_resize'

# 确保输出文件夹存在，如果不存在则创建
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历输入文件夹中的所有png文件
for filename in os.listdir(input_folder):
    if filename.endswith(".png"):
        input_path = os.path.join(input_folder, filename)
        # 打开图像文件
        img = Image.open(input_path)
        # 调整图像大小为256x256
        resized_img = img.resize((64, 64))
        # 构造输出文件的路径
        output_path = os.path.join(output_folder, filename)
        # 保存调整大小后的图像
        resized_img.save(output_path)
        print(f'Resized {input_path} to {output_path}')

为了比较不同情况下的数据集训练出来的结果，我把cubicasa5k分为了cubicasa5k-simple和cubicasa5k-large两个版本

1
2
3
4
5
6
7
8
├── cubicasa5k-simple
│       ├── colorful
│       ├── high_quality
│ 
├── cubicasa5k-large
│       ├── colorful
│       ├── high_quality
│       ├── high_quality_architectural

ok，终于处理完了数据集，开始训练把~

对cubicasa5k的训练

刚才分为了两个版本，那我们分开训练

对cubicasa5k-simple的训练

image-20231205152736322

不太清楚要训练多少轮，因为不同的数据集不同的轮次训练出来的效果是不一样的，没有一个固定的标准。

我这里训练了600个epoch，其实从训练中反馈的信息来看，训练到200轮的时候就已经收敛的差不多了：

image-20231205151402675

那么最终训练完600轮后效果如下：

image-20231205151512710

对cubicasa5k-large的训练

image-20231205152724359

cubicasa5k-large和cubicasa5k-simple的区别就是，cubicasa5k-large包含了更多复杂结构户型图数据，所以训练起来是要比simple版本难得，实际实验看出，在训练到200轮的时候，依旧部分收敛，但是还是存在很明显的高斯噪声分布：

image-20231205151701628

最终训练完600轮的效果如下：

image-20231205151725880

可以看出，大部分是采样的比较好的，偶尔有一两个噪声分布比较明显。

对RPLAN的训练

image-20231205152752428

RPLAN这个数据集比较大，有八万多张图片，所以训练起来就比较难，**在 2*3090（24G） GPU的服务器上，batchsize=8，训练一轮要9分40秒，整体的训练代价是比较大的，所以这里我总共训练了200轮**。

第100轮时的效果：

image-20231205152108892

第200轮时的效果：

image-20231205152202517

可以看到收敛的速度还是比较快的，因为这个数据比较简单，没有太复杂分布。

再者说，我感觉这个数据集主要的任务是：Layout Generation，可能用在这次设计中有点不太合适。

对自己构建的数据集进行训练

image-20231205152807452

之前也提到了，我从二手房网站上爬了一些数据，构建了自己的数据集，经过人工清洗之后，可用的数据量大概3000多张：

image-20231205152634326

针对这个数据集，我也训练了600轮，看看效果吧：

200轮时效果：

image-20231205152905617

可以看到由于色彩较多，原始数据集的分布比较复杂，所以200轮的采样结果不是很理想，这里出现绿色是因为有些户型图数据中有草坪：

image-20231205153037085

所以说针对这个数据集，多训练几轮是很有必要的，600轮的情况如下：

image-20231205153133228

可以看出，600轮的时候就差不多了，由于实验室还有其他的工作要做，所以这次的训练就到此为止了。

测试

对cubicasa5k-Simple-DM的测试

对应参数调好：

image-20231205153432097

生成指令：

image-20231205153632303

输出结果：

image-20231205153912089

还是比较ok的

对cubicasa5k-Large-DM的测试

对应参数调好：

image-20231205153947293

生成指令：

image-20231205154025304

输出结果：

image-20231205154046851

也算凑合

对RPLAN-DM的测试

对应参数调好：

image-20231205154131446

生成指令：

image-20231205154241236

输出结果：

image-20231205154308035

有一张崩了，剩下还好

对Myself-DM的测试

对应参数调好：

image-20231205154339712

生成指令：

image-20231205154416779

输出结果：

image-20231205154438642

ok，可以看到也能生成对应的图像，只不过噪声比较多，后续有机会再补上剩下的训练吧。

总结

本次我训练了四个版本的Diffusion Model for Floor Plan datasets，分别是：

  • cubicasa5k-Simple-DM
  • cubicasa5k-Large-DM
  • RPLAN-DM
  • Myself-DM（自己从网上爬取的数据集）

后续该设计有新的实现也会同步更新，在设计验收后，我会将代码上传github。

扫一扫，分享到微信

tag:

    缺失模块。
    1、请确保node版本大于6.2
    2、在博客根目录（注意不是yilia-plus根目录）执行以下命令：
    npm i hexo-generator-json-content --save

    3、在根目录_config.yml里添加配置： 

      jsonContent:
    meta: false
    pages: false
    posts:
      title: true
      date: true
      path: true
      text: false
      raw: false
      content: false
      slug: false
      updated: false
      comments: false
      link: false
      permalink: false
      excerpt: false
      categories: false
      tags: true

计算机小学徒，慢慢成长