1D Diffusers not behaving as expected, am i retarded?

Beginners
1

This is my first Post on the Forum, as I’m usually able to solve my issues on my own. But in this very simple Diffusion training of a trivial 1D dataset, the model just cant seem to figure it out. This leads me to think that I’m fundamentally misunderstanding something in the 1D diffusers implementation. Or perhaps it is not supported how i imagined it to be.

I have trained 2D models without issues for reference, with close to identical code.

The following code is the entire program, which should return a graph at each 50 epochs (or whatever you set it to). I tried to set the epoch to something fairly high, to verify that it is not the training time. The returned graphs range between [400, -400], and does not seem to improve with time training. The training has a loss at 0.7 at the first epoch and then goes down to 0.51 and stays there. Which i believe is also wrong.

code:

import torch
from diffusers import DDPMScheduler, UNet1DModel
from torch.utils.data import Dataset, DataLoader
from matplotlib import pyplot as plt
import numpy as np

class Synthetic1DData(Dataset):
    def __init__(self, num_samples=1000, seq_len=64):  
        self.data = torch.ones(num_samples, 1, seq_len)
        self.data[:, :, 32:] = 0
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]
model = UNet1DModel(
    sample_size=64,
    in_channels=1,
    out_channels=1,
    layers_per_block=1,
    block_out_channels=(32, 64),
    down_block_types=("DownBlock1D", "DownBlock1D"),
    up_block_types=("UpBlock1D", "UpBlock1D"),
)
noise_scheduler = DDPMScheduler(
    num_train_timesteps=1000,
    clip_sample=False
)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

dataset = Synthetic1DData(seq_len=64)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

model.train()
for epoch in range(1000):
    total_loss = 0
    num_batches = 0
    for batch in dataloader:
        noise = torch.randn_like(batch)
        
        timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (batch.shape[0],))
        noisy = noise_scheduler.add_noise(batch.to(device), noise.to(device), timesteps.to(device))
        pred = model(noisy.to(device), timesteps.to(device)).sample

        loss = torch.nn.functional.mse_loss(pred.to(device), noise.to(device))
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        num_batches += 1

    #Logging
    if epoch % 50 == 0:
        avg_loss = total_loss / num_batches
        print(f"Epoch {epoch}, Average Loss: {avg_loss:.4f}")

        model.eval()
        with torch.no_grad():
            sample = torch.randn(1, 1, 64).to(device)
            for t in noise_scheduler.timesteps:
                residual = model(sample, t).sample
                sample = noise_scheduler.step(residual, t, sample).prev_sample
            
            plt.figure(figsize=(10, 4))
            plt.plot(np.arange(64), sample[0].cpu().squeeze().numpy())
            plt.title(f'Epoch {epoch}')
            plt.show()
        model.train()
1 Like
2

I think something wrong in UNet1DModel…