colab: 所有可用的 RAM 皆已用盡，因此你的工作階段已停止運作

2025-02-282025-02-28

在 Colab 中，GPU 內存 (RAM) 使用過高時，會導致工作階段停止運作。為了減少內存消耗，你可以考慮以下幾種方法：

1. 減少 `batch_size`:

減少每次訓練步驟中處理的樣本數，可以有效降低 GPU 內存的使用量。

training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=1,  # 減少 batch size
    gradient_accumulation_steps=4,  # 保持積累梯度，讓訓練不會過快停止
    learning_rate=1e-4,
    max_steps=50,
    save_strategy="steps",
    save_steps=10,
    report_to="none",
    remove_unused_columns=False
)

2. 使用 `gradient_accumulation_steps`:

這樣做的目的是將小批量的梯度累積起來，並在較少的步驟中執行反向傳播。通過這種方式，你可以有效地模擬較大的 batch_size，但實際上使用較小的內存。

3. 減少模型大小:

如果你的模型比較大，可以考慮選擇較小的模型（例如，使用 stable-diffusion-2 的變種版本）來減少內存消耗。

4. 使用半精度浮點數 (FP16):

將模型和數據轉換為半精度浮點數（float16）可以顯著減少內存佔用。

在你的程式碼中，模型的載入和數據都已經被設置為使用 float16：

unet = UNet2DConditionModel.from_pretrained(
    PRETRAINED_MODEL_NAME,
    subfolder="unet",
    torch_dtype=torch.float16  # 使用半精度
)

5. 減少圖片解析度:

高解析度圖片需要更多的內存。你可以通過調整圖片大小來減少 GPU 內存的消耗。

def transform_gray_to_rgb(example):
    img = example["image"]
    img = img.convert("RGB")
    
    # 在訓練前將圖片調整為較低解析度
    img = img.resize((256, 256))  # 將圖像調整為 256x256 或更小
    img = transforms.ToTensor()(img)
    
    fourth_channel = torch.ones(img.shape[1], img.shape[2], dtype=img.dtype, device=img.device)
    img = torch.cat([img, fourth_channel.unsqueeze(0)], dim=0)
    
    example["pixel_values"] = img
    return example

6. 使用更小的 LoRA Rank:

減少 LoRA 的 r 參數值可以減少內存占用。例如，將 r=8 改為 r=4 或 r=2。

lora_config = LoraConfig(
    r=4,  # 減少 LoRA Rank
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=[
        "attn1.to_q", "attn1.to_k", "attn1.to_v", "attn1.to_out",  
        "attn2.to_q", "attn2.to_k", "attn2.to_v", "attn2.to_out"  
    ],
    bias="none",
)

7. 卸載不必要的模組:

在進行訓練過程中，若有其他不需要的模型或變量佔用了大量內存，記得適時清理。可以使用 torch.cuda.empty_cache() 來釋放未使用的顯存。

import torch

# 在每次訓練步驟後釋放未使用的顯存
torch.cuda.empty_cache()

8. 使用更小的隱藏層大小或序列長度:

如果模型的隱藏層（例如 Transformer 層）很大，將其設置為較小的大小也可以節省內存。

這些方法應該能幫助你減少 Colab 上 GPU 內存的使用，避免訓練過程中 GPU 內存耗盡的情況。

Max的程式語言筆記