# UNet1DModel not converging on single batch
**URL:**
**Author:**  [@StableDiffuser317](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FStableDiffuser317)%5C
**Post date:** [April 25, 2023, 8:02am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Funet1dmodel-not-converging-on-single-batch%2F37656%2F1 "2023-04-25T08:02:27Z")
Hello,
Iâm currently trying to implement a latent generative diffuser using a custom-trained VAE which maps 3x256x256 to 1x256 and a UNet1DModel with the following parameters :
UNet1DModel(in\_channels = 1, out\_channels=1,down\_block\_types =
(âDownBlock1Dâ, âDownBlock1Dâ, âDownBlock1Dâ,âAttnDownBlock1Dâ,âDownBlock1Dâ),
up\_block\_types = (âUpBlock1Dâ,âUpBlock1Dâ, âUpBlock1Dâ, âAttnUpBlock1Dâ, âUpBlock1Dâ),
block\_out\_channels = (32, 64, 128, 256, 512), time\_embedding\_type = âfourierâ,layers\_per\_block = 2, act\_fn = âsiluâ)
Training is however pretty tough as I canât reach a small enough training loss, I need something of the order of 1e-4.
Here is the training code :
"
# Lets train
from accelerate import Accelerator
from diffusers.hub\_utils import init\_git\_repo, push\_to\_hub
import os
def train\_loop(config, model, vae,noise\_scheduler, optimizer, train\_dataloader, val\_dataloader, lr\_scheduler):
# Initialize accelerator and tensorboard logging
accelerator = Accelerator(
mixed\_precision=config.mixed\_precision,
gradient\_accumulation\_steps=config.gradient\_accumulation\_steps,
log\_with=âtensorboardâ,
logging\_dir=os.path.join(config.output\_dir, âlogsâ)
)
if accelerator.is\_main\_process:
accelerator.init\_trackers(âtrain\_exampleâ)
model, vae, optimizer, train\_dataloader, val\_dataloader,lr\_scheduler = accelerator.prepare(
model, vae, optimizer, train\_dataloader, val\_dataloader,lr\_scheduler
)
```auto
global_step = 0
# Now you train the model
epoch_wise_val_loss = []
for epoch in range(config.num_epochs):
batch = next(iter(train_dataloader))
if epoch%500 == 0:
print(f"Epoch : {epoch}",flush=True)
with torch.no_grad():
batch = batch.to(torch_device)
mu, log_var = vae.encode(batch)
clean_images = vae.reparameterize(mu, log_var).unsqueeze(1)
# Sample noise to add to the images
noise = torch.randn(clean_images.shape).to(clean_images.device)
bs = clean_images.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bs,), device=clean_images.device).long()
# Add noise to the clean images according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
with accelerator.accumulate(model):
# Predict the noise residual
noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
loss = F.mse_loss(noise_pred, noise)
accelerator.backward(loss)
accelerator.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
accelerator.log(logs, step=global_step)
global_step += 1
# After each epoch you optionally sample some demo images with evaluate() and save the model
if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
unwrapped_model = accelerator.unwrap_model(model)
#unwrapped_model.save_pretrained(save_directory=config.output_dir+"/epoch_"+str(epoch+1))
pipeline_trained = LDMPipeline(vae=vae,unet=unwrapped_model,scheduler=noise_scheduler,input_device=torch_device)
evaluate(config,epoch+1,pipeline_trained)
```
Im using a cosine scheduler with a lr of 2e-5 (Iâve tried e-4 to e-8) and I donât need to normalize the data as they already have mean=0 and std=1 after being mapped by the VAE.
The training loss is :
Have you got any idea on the source of the problem?
Thanks
---
**Author:**  [@drimeF0](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FdrimeF0)%5C
**Post date:** [June 2, 2023, 4:04pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Funet1dmodel-not-converging-on-single-batch%2F37656%2F2 "2023-06-02T16:04:59Z")
hi i have the same problem but with unet2D
code here: [Google Colab](https://colab.research.google.com/drive/1AjH6ujmLIsU-4hP0bX0jtaFv_WGHpu2o)
I tried a lot of things (decreased the learning\_rate, set different prediction\_types, changed the architecture of the neural network)
---
**Author:**  [@drimeF0](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FdrimeF0)%5C
**Post date:** [June 2, 2023, 4:17pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Funet1dmodel-not-converging-on-single-batch%2F37656%2F3 "2023-06-02T16:17:47Z")
maybe I found a solution to the problem : [LDM Training ¡ Issue #2164 ¡ huggingface/diffusers ¡ GitHub](https://github.com/huggingface/diffusers/issues/2164#issuecomment-1424204308)
---
**Author:**  [@StableDiffuser317](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FStableDiffuser317)%5C
**Post date:** [June 19, 2023, 1:01pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Funet1dmodel-not-converging-on-single-batch%2F37656%2F4 "2023-06-19T13:01:09Z")
did it work ? can send u my code if u need
---
**Author:**  [@ElcaidaProjects](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FElcaidaProjects)%5C
**Post date:** [May 10, 2026, 11:13am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Funet1dmodel-not-converging-on-single-batch%2F37656%2F5 "2026-05-10T11:13:46Z")
this is the fix to your exact architecture: UNet1DModel(in\_channels = 1, out\_channels=1,extra\_in\_channels=64, down\_block\_types =
(âDownBlock1DNoSkipâ, âDownBlock1Dâ, âDownBlock1Dâ,âAttnDownBlock1Dâ,âDownBlock1Dâ),
up\_block\_types = (âUpBlock1Dâ,âUpBlock1Dâ, âUpBlock1Dâ, âAttnUpBlock1Dâ, âUpBlock1DNoSkipâ),
block\_out\_channels = (32, 64, 128, 256, 512), time\_embedding\_type = âfourierâ,layers\_per\_block = 2, act\_fn = âsiluâ) â i opened an issue on this specific quirk on github as well but it was quickly closed, claiming to be a feature and not a bug