the stability of training(a collapse loss)

[sunsq-blue]

Thanks for the great work. I try to train the ldm model on ImageNet with 8 V100, but get a bad result.I found that loss was normal at first, but soon collapsed:

and the sampled image are all noise at 5000 steps:

How can I solve this problem, thank you very much

[sunsq-blue]

update: The configuration file has hardly been modified。

model: base_learning_rate: 1.0e-06 target: ldm.models.diffusion.ddpm.LatentDiffusion params: linear_start: 0.0015 linear_end: 0.0195 num_timesteps_cond: 1 log_every_t: 200 timesteps: 1000 first_stage_key: image cond_stage_key: class_label image_size: 32 channels: 4 cond_stage_trainable: true conditioning_key: crossattn monitor: val/loss_simple_ema unet_config: target: ldm.modules.diffusionmodules.openaimodel.UNetModel params: image_size: 32 in_channels: 4 out_channels: 4 model_channels: 256 attention_resolutions:

note: this isn\t actually the resolution but

    # the downsampling factor, i.e. this corresnponds to
    # attention on spatial resolution 8,16,32, as the
    # spatial reolution of the latents is 32 for f8
    - 4
    - 2
    - 1
    num_res_blocks: 2
    channel_mult:
    - 1
    - 2
    - 4
    num_head_channels: 32
    use_spatial_transformer: true
    transformer_depth: 1
    context_dim: 512
first_stage_config:
  target: ldm.models.autoencoder.VQModelInterface
  params:
    embed_dim: 4
    n_embed: 16384
    ckpt_path: models/first_stage_models/vq-f8/model.ckpt
    ddconfig:
      double_z: false
      z_channels: 4
      resolution: 256
      in_channels: 3
      out_ch: 3
      ch: 128
      ch_mult:
      - 1
      - 2
      - 2
      - 4
      num_res_blocks: 2
      attn_resolutions:
      - 32
      dropout: 0.0
    lossconfig:
      target: torch.nn.Identity
cond_stage_config:
  target: ldm.modules.encoders.modules.ClassEmbedder
  params:
    embed_dim: 512
    key: class_label

data: target: main.DataModuleFromConfig params: batch_size: 64 num_workers: 12 wrap: false train: target: ldm.data.imagenet.ImageNetTrain params: data_root: /data/imagenet/ImageNet/data config: size: 256 validation: target: ldm.data.imagenet.ImageNetValidation params: data_root: /data/imagenet/ImageNet/data config: size: 256

lightning: callbacks: image_logger: target: main.ImageLogger params: batch_frequency: 5000 max_images: 8 increase_log_steps: False

trainer: benchmark: True

[sunsq-blue]

update:
When I train on single GPU , the loss collapse seem not to happen again. And another solution may change the AdamW to Adam. I hope others will share their solutions，thanks.

[CrossLee1]

@sunsq-blue I also met the same problem with you. Did you succeed to make it work using the above solutions? Thanks

[Ir1d]

@CrossLee1 @sunsq-blue I figure out that you shouldn't scale the lr too much. You should use --scale_lr False and it works well. In lightning the effective batch size is batch_size n_gpu n_device. In code they scale the lr to be bs * n_gpu

[wtliao]

@Ir1d Thanks a lot! I was suffering the same problem for several days. Finally solved with your suggestion!