stable diffusion中dreambooth模型训练_runtimeerror: cuda out of memory. tried to allocat

简介

一、参考材料

1、Github项目网址
2、huggingface网址
3、wandb网址

二、运行备注

1、背景介绍
由于之前没有运行Github项目的经验，这个项目是我在Github尝试进行的第一个项目，训练过程中遇到了很多报错，花费了大量时间熟悉服务器、wsl的用法、vs code的连接方式等。在初期我希望能够找到这个项目环境配置的相关教学资料，发现很多材料都不是很熟悉，对于我这样的新手小白来说跟进十分困难，因此在本篇文章中将尽量详细地进行项目配置介绍，希望能够帮助到以后使用此项目上手的同学。

2、模型训练信息
（1）模型在服务器（Linux系统）上训练。由于我使用的是windows系统，使用wsl虚拟机连接服务器
注：需要提前配置wsl虚拟环境（如安装Ubuntu等）
（2）编译器为vs code，具体连接方式可见博客文章《vs code通过ssh连接服务器报错》

模型训练步骤

注：在进行模型训练之前切记先阅读项目目录下的readme文件

一、连接服务器

1、方式一：如果下载了Ubuntu，直接在开机处搜索栏搜索Ubuntu打开终端输入以下代码，按照提示输入登录密码即可连接到远程服务器。

ssh user_name@server_IP
1

2、方式二：使用powershell连接远程服务器。按下Win + R打开“运行”对话框，输入Powershell打开Powershell界面，输入wsl，加载完毕后再按照1中步骤即可连接远程服务器。

二、虚拟环境配置

1、阅读stable diffusion首页readme文件可知，推荐用PyPI或conda在虚拟环境中安装diffusers。阅读dreambooth中readme文件可知，主要步骤为创建虚拟环境→安装diffusers→安装依赖→加速器配置→模型训练。
注：在此步骤之前，确保服务器已经安装conda或pip。在以下介绍中涉及虚拟环境创建或激活的指令均以conda为例。

2、导航到服务器中目标文件夹位置，创建虚拟环境

cd dest_location # dest_location替换为目标文件夹地址
conda create -n env_name python==3.10 # 此处填写所需python版本即可
1
2

3、导航到目标文件夹位置（最好与创建的虚拟环境在同一文件夹下），安装diffusers，输入以下指令安装diffusers
注：不要开代理或梯子，避免出现报错

conda install -c conda-forge diffusers # 安装diffusers
1

4、安装环境依赖，导航到diffusers母文件夹，输入以下指令。

cd diffusers
pip install -e .
1
2

5、进入diffusers/example/dreambooth文件夹，执行以下指令。
注：如果出现有包未安装的警告，则按照提示进入相应网址下载安装相应安装包

pip install -r requirements_sdxl.txt
accelerate config
1
2

输入accelerate config后出现若干选择题，一般情况下选择默认即可（默认使用1张GPU卡）。若希望同时使用多张GPU卡进行模型训练，则可按照以下文章进行显卡配置（多卡训练模型accelerate配置）

6、在diffusers/examples/dreambooth路径下新建一个download.py文件，内容如下：

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir, repo_type="dataset",
    ignore_patterns=".gitattributes",
)
1
2
3
4
5
6
7
8

在同一路径下新建一个model_test.sh文件（训练模型指令），内容如下：

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0" 
export INSTANCE_DIR="dog" # 存放学习图片文件夹
export OUTPUT_DIR="lora-trained-xl" # 存放输出结果图片文件夹
export VAE_PATH="madebyollin/sdxl-vae-fp16-fix"

accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --pretrained_vae_model_name_or_path=$VAE_PATH \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \ # 输出图片指令
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

由model_test.sh文件可知，需要在model_test.sh统一母文件夹下新建一个存放学习图片的文件夹（dog）以及一个存放输出结果的文件夹（lora-trained-xl）。

7、（已连接服务器的情况下）打开创建的虚拟环境，指令如下所示：

source env_name/bin/activate # 注意此时需要先导航到虚拟环境所在母文件夹下
1

连接到创建的虚拟环境后，终端显示如下所示（我创建的环境名称为dreambooth_env）：

8、在虚拟环境中运行download.py文件，下载所需的学习文件。指令为：

python download.py
1

注：在此过程中可能需要获取huggingface官网的API，按照提示在hf官网进行注册即可。

9、在虚拟环境中进行模型训练，指令如下

chmod +x model_test.sh
./model_test.sh
1
2

在运行训练文件之前确保虚拟环境中已经下载了wandb文件，指令如下：

pip install wandb
1

在此过程中可能需要获取wandb的API，在wandb官网进行注册即可。

10、在服务器及GPU有剩余内存的情况下，应该是能够正常进行模型训练的。如果未能成功运行，出现以下报错可以按照以下方式进行解决。

• GPU内存不足（CUDA out of memory. Tried to allocate xxx MiB. GPU）
  • 解决方式：在服务器中查看GPU内存使用情况，指令为：

    nvidia-smi
    1

    输出表头包含每栏内容主题，如GPU名称、风扇转速、GPU运行状态、内存使用情况等，具体内容可见查看服务器显卡状态文章。
    如果有GPU剩余较多内存，则说明默认使用的芯片内存不足，但系统并未能调用其余芯片进行模型训练。此时可以对训练模型使用到的GPU进行指定。指定方式为在model_test.sh文件中修改（accelerate launch train_dreambooth_lora_sdxl.py \）语句为以下内容：

    accelerate launch CUDA_VISIBLE_DEVICES=xx train_dramvooth_lora_sdxl.py \
    # 其中xx替代为所需指定的GPU卡号，通过nvidia-smi可以查看GPU卡号
     ```
    1
    2
    3

• 找不到学习图片报错，具体如下所示：

Traceback (most recent call last):
  File ".../diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 1984, in <module>
    main(args)
  File ".../diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 1415, in main
    train_dataset = DreamBoothDataset(
  File ".../diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 780, in __init__
    instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())]
  File ".../diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py", line 780, in <listcomp>
    instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())]
  File ".../dreambooth_env/lib/python3.10/site-packages/PIL/Image.py", line 3431, in open
    fp = builtins.open(filename, "rb")
IsADirectoryError: [Errno 21] Is a directory: '.../diffusers/examples/dreambooth/dog/.huggingface'
1
2
3
4
5
6
7
8
9
10
11
12

上述报错说明尝试打开图片文件时给出的路径下为文件夹而非图片文件，具体解决方式为修改train_dreambooth_lora_sdxl.py文件第780行左右的代码，修改方式如下：

# 原代码
instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir())]

# 过滤掉目录，只保留文件
instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir()) if path.is_file()]

# 或者只处理指定格式的图像文件
instance_images = [Image.open(path) for path in list(Path(instance_data_root).iterdir()) if path.is_file() and path.suffix in ['.jpg', '.jpeg', '.png']]
1
2
3
4
5
6
7
8

• 以下报错可以忽略

subprocess.CalledProcessError: Command '['./dreambooth_env/bin/python', 'train_dreambooth_lora_sdxl.py', '--pretrained_moth=stabilityai/stable-diffusion-xl-base-1.0', '--instance_data_dir=dog', '--pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16ut_dir=lora-trained-xl', '--mixed_precision=fp16', '--instance_prompt=a photo of sks dog', '--resolution=1024', '--train_batch_size=1',ccumulation_steps=4', '--learning_rate=1e-4', '--report_to=wandb', '--lr_scheduler=constant', '--lr_warmup_steps=0', '--max_train_stepsdation_prompt=A photo of sks dog in a bucket', '--validation_epochs=25', '--seed=0', '--push_to_hub']' returned non-zero exit status 1.
1

注：上述文件注意存放在同一文件夹下。如果仍有问题可以在Github项目中的issue栏目进行搜索 or 提问 / 询问ChatGPT / 浏览器搜索，Github上issue位置如下所示。