class MobileNetModule(nn.Module):
def init(self, dim, time_dim, cond_dim, expand_ratio=6):
super().init()
hidden_dim = dim*expand_ratio
self.expand_conv = nn.Sequential(
nn.Conv2d(dim, hidden_dim, kernel_size=1, bias=False),
nn.BatchNorm2d(hidden_dim),
nn.ReLU6()
)
self.dwise_conv = nn.Sequential(
nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, groups=hidden_dim, padding=1, bias=False),
nn.BatchNorm2d(hidden_dim),
nn.ReLU6()
)
self.proj_conv = nn.Sequential(
nn.Conv2d(hidden_dim, dim, kernel_size=1, bias=False),
nn.BatchNorm2d(dim),
)
self.t_emb = nn.Sequential(
nn.Linear(time_dim, dim),
nn.ReLU()
)
self.cond_proj = nn.Linear(cond_dim, dim)

def forward(self, x, t, cond):
t = self.t_emb(t).unsqueeze(2).unsqueeze(3)
if cond is not None:
cond = self.cond_proj(cond).unsqueeze(2).unsqueeze(3)
t = t + cond

out = self.expand_conv(x)
out = self.dwise_conv(out)
out = self.proj_conv(out)
return out + t


Modified HalfUNetBlock for latent space
class HalfUNetBlock(nn.Module):
def init(self, block: nn.Module, dim, time_dim, cond_dim):
super(HalfUNetBlock, self).init()
self.block1 = block(dim, time_dim=time_dim, cond_dim=cond_dim)
self.block2 = block(dim, time_dim=time_dim, cond_dim=cond_dim)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

def forward(self, x, t, cond):
x = self.block1(x, t, cond)
x = self.block2(x, t, cond)
p = self.pool(x)
return x, p
Modified HalfUNetModel for latent space
class HalfUNetModel(nn.Module):
def init(self, block: nn.Module, dim, depth, time_dim=128, cond_dim=1000):
super(HalfUNetModel, self).init()
self.depth = depth
self.blocks = nn.ModuleList()
for _ in range(depth):
self.blocks.append(HalfUNetBlock(block, dim=dim, time_dim=time_dim, cond_dim=cond_dim))

def forward(self, x, t, cond):
skip_connections = []
for block in self.blocks:
x, p = block(x, t, cond)
skip_connections.append(x)
x = p
out = torch.zeros_like(skip_connections[0])
for i, skip in enumerate(skip_connections):
out += F.interpolate(skip, size=out.shape[2:])
return out

Modified NestedHalfUNet for latent space

class NestedHalfUNet(nn.Module):
def init(self, block: nn.Module, ch_in, ch_out, depth, dim, time_dim, num_classes=1000):
super(NestedHalfUNet, self).init()
self.time_dim = time_dim
self.init_conv = nn.Conv2d(ch_in, dim, kernel_size=1)
self.down = nn.ModuleList()
for i in range(depth):
self.down.append(HalfUNetModel(block, depth=depth-i, dim=dim, time_dim=time_dim, cond_dim=num_classes))
self.final_conv = nn.Conv2d(dim, ch_out, kernel_size=1)
self.buff = self.register_buffer('inv_freq', 1.0 / (
10000 ** (torch.arange(0, time_dim, 2).float() / time_dim)
))

def pos_encoding(self, t, channels):
pos_enc_a = torch.sin(t.repeat(1, channels // 2) * self.inv_freq)
pos_enc_b = torch.cos(t.repeat(1, channels // 2) * self.inv_freq)
pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)
return pos_enc

def forward(self, x, t, label=None):
t = t.unsqueeze(-1).type(torch.float)
t = self.pos_encoding(t, channels=self.time_dim)
x = self.init_conv(x)
skip_connections = []
for down in self.down:
x = down(x, t, label)
skip_connections.append(x)
x = F.interpolate(x, scale_factor=0.5)
out = torch.zeros_like(skip_connections[0])
for i, skip in enumerate(skip_connections):
out += F.interpolate(skip, size=out.shape[2:])
return self.final_conv(out)

Training setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

lr = 3e-4
num_epochs = 20
snr_gamma = 5.0
parameterization = 'v'

optimizer = optim.AdamW(model.parameters(), lr=lr)
scaler = amp.GradScaler(device=device)
mse = F.mse_loss
diffusion = Diffusion(cosine_noise_sched=True, conditional=False, img_size=latent_size, channels=4, device=device, snr_gamma=snr_gamma)

Training function

def train(model, latents, labels):
t = diffusion.sample_timesteps(latents.shape[0]).to(device)

x_t, noise = diffusion.noise_images(latents, t)

if diffusion.parameterization == "eps":
    target = noise
elif diffusion.parameterization == "x0":
    target = latents
elif diffusion.parameterization == "v":
    target = diffusion.get_v(latents, noise, t)

with amp.autocast(device_type=device.type):
    model_out = model(x_t, t, label=labels)
    loss = mse(model_out, target, reduction='none' if diffusion.snr_gamma else 'mean')

if diffusion.snr_gamma:
    snr = diffusion.compute_snr(t)
    base_weight = (torch.stack([snr, diffusion.snr_gamma * torch.ones_like(t)], dim=1).min(dim=1)[0] / snr)
    base_weight[snr == 0] = 1.0
    loss = loss.mean(dim=list(range(1, len(loss.shape)))) * base_weight
    loss = loss.mean()

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

optimizer.zero_grad()
return loss

Training loop

epoch_pbar = tqdm(range(num_epochs))
curr_step = 0
total_steps = num_epochs*len(dataloader)
losses = []

for epoch in epoch_pbar:
loader_pbar = tqdm(dataloader, leave=False, total=len(dataloader))
for i, batch in enumerate(loader_pbar):
latents = batch["vae_output"].reshape(-1, 4, 32, 32).to(device).float()
labels = F.one_hot(torch.as_tensor([int(lbl) for lbl in batch["label"]]), num_classes=1000).float().to(device)
if np.random.uniform() < 0.5: labels = torch.zeros_like(labels).float().to(device)

    loss = train(model, latents, labels)
    

    if curr_step%5000==0:
        sample_label = F.one_hot(torch.tensor([3]), num_classes=1000).float().to(device)
        latent = diffusion.sample(model, 1, labels=sample_label)
        img = (vae.decode(latent.to(torch.float16))['sample'].squeeze().permute(1,2,0).detach().cpu().numpy().clip(-1, 1)/2+0.5).astype(np.float32)
        plt.imshow(img)
        plt.show()

    loader_pbar.set_postfix(MSE=loss.item())
    curr_step += 1
    losses.append(loss.item())

streaming.base.util.clean_stale_shared_memory()


here is model architecture and training loop of latent diffusion image generator model. I have a goal to make smallest image generator for ImageNet, but current model doesn't learn anything. How to improve model without making it bigger (22M params now)? Think about better base block (now MobileNetModule) or "macro architecture" (now half unet)