Problem of training and inference on the COCO dataset

[cqugzt]

I tried to perform training and inference on the COCO dataset.

But the trained model cannot generate correct images based on the text.

• A sample:

caption : A woman is standing outside with her head on a bat.

• LR and Loss:

What should I do next?

I would appreciate your advice very much！

[cqugzt]

Here is my code:

from transfusion_pytorch import Transfusion, print_modality_sample
from diffusers.models import AutoencoderKL  

import json  
from PIL import Image  
from torch.utils.data import Dataset  
import torchvision.transforms as T  
from itertools import cycle
from torch.optim.lr_scheduler import LambdaLR

import torch
import math
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import save_image
from transformers import AutoTokenizer

import os  
from pathlib import Path  
from torch.optim import AdamW
from tqdm import tqdm  

# === Hyperparameters Configuration ===
DO_TEST = True
TRAIN_FILE_PATH = '/root/autodl-tmp/.autodl/transfusion-pytorch-main/src/dataset/val_dataset_split/train_dataset.json'
TEST_FILE_PATH = '/root/autodl-tmp/.autodl/transfusion-pytorch-main/src/dataset/val_dataset_split/test_dataset.json'
VAL_FILE_PATH = '/root/autodl-tmp/.autodl/transfusion-pytorch-main/src/dataset/val_dataset_split/val_dataset.json'
IMAGE_DIR = '/root/autodl-fs/val2017'

# === Model Configuration ===

# Load pretrained VAE for handling images  
vae = AutoencoderKL.from_pretrained("/root/autodl-fs/pre-models/stable-diffusion-v1-4",subfolder="vae")

# Define the encoder using a pretrained VAE model  
class Encoder(nn.Module):  
    def __init__(self, vae):  
        super().__init__()  
        self.vae = vae  # Store the pretrained VAE (e.g., from Stable Diffusion)

    def forward(self, image):  
        # Disable gradient computation (faster inference, less memory)
        with torch.no_grad():
            # Normalize input image from [0, 1] to [-1, 1] before encoding
            latent = self.vae.encode(image * 2 - 1)
        
        # Sample from the latent distribution and scale it (Stable Diffusion uses 0.18215 as a scaling factor)
        return 0.18215 * latent.latent_dist.sample()


# Define the decoder using the same pretrained VAE model  
class Decoder(nn.Module):  
    def __init__(self, vae):  
        super().__init__()  
        self.vae = vae  # Use the same pretrained VAE

    def forward(self, latents):  
        # Rescale latent back to the original range expected by the VAE decoder
        latents = (1 / 0.18215) * latents

        # Decode the latent into an image using the VAE decoder
        with torch.no_grad():
            image = self.vae.decode(latents).sample  # Sample the decoded output

        # Convert the image back from [-1, 1] to [0, 1] and clip values to valid range
        return (image / 2 + 0.5).clamp(0, 1)


TOKENIZER = AutoTokenizer.from_pretrained("/root/autodl-fs/pre-models/llama-2")

# Helper functions for text encoding/decoding  
def encode_tokens(text):  
    return torch.tensor(TOKENIZER.encode(text), dtype=torch.int) 
  
def decode_tokens(tokens):  
    return "".join(TOKENIZER.decode(tokens, skip_special_tokens=True))  


# Configure for Transfusion model
model = Transfusion(  
    num_text_tokens=TOKENIZER.vocab_size,  # Increased for text vocabulary
    dim_latent=4,         # VAE latent dimension
    channel_first_latent=True,
    modality_default_shape=(32, 32),  # Default shape for COCO latents
    modality_encoder=Encoder(vae),
    modality_decoder=Decoder(vae),
    # pre_post_transformer_enc_dec=(  
    #     nn.Conv2d(4, 384, 3, 2, 1),
    #     nn.ConvTranspose2d(384, 4, 3, 2, 1, output_padding=1),
    # ),  
    add_pos_emb=True,
    modality_num_dim=2,
    transformer=dict(  
        dim=384,           # Increased model dimension  
        depth=8,           # Deeper model  
        heads=8,           # attention heads
        dim_head=64
    )

).cuda()  
  
# Create EMA model for stable sampling  
ema_model = model.create_ema(0.9)

# === Dataset Configuration ===

# Define dataset class for COCO
class COCODataset(Dataset):  
    def __init__(self, json_file, image_dir, image_size=256):  
        # Load annotations from JSON file  
        with open(json_file, 'r') as f:  
            self.annotations = json.load(f)  
          
        self.image_dir = image_dir  
          
        # Define image transformations  
        self.transform = T.Compose([  
            T.Resize((image_size, image_size)),  
            T.PILToTensor(),  
            T.Lambda(lambda t: t / 255.)  
        ])  
      
    def __len__(self):  
        return len(self.annotations)
      
    def __getitem__(self, idx):  
        # Access the specific keys from your JSON structure  
        annotation = self.annotations[idx]  
        caption = annotation["caption"]  
        file_name = annotation["file_name"]  
          
        # Load the image  
        image_path = os.path.join(self.image_dir, file_name)  
        image = Image.open(image_path).convert("RGB")  
          
        # Apply transformations  
        image_tensor = self.transform(image)  
        
        # For text processing, you might need to tokenize or encode the caption
        text_tensor = encode_tokens(caption)

        return text_tensor, image_tensor


# Create dataloader  
train_dataset = COCODataset(
    json_file=TRAIN_FILE_PATH,   
    image_dir=IMAGE_DIR,  
    image_size=256  
) 

dataloader = model.create_dataloader(train_dataset, batch_size=8, shuffle=True)  
iter_dl = cycle(dataloader)  # Create infinite iterator

# === Training Loop ===

# Initialize TensorBoard writer
writer = SummaryWriter('./logs')

# Create directories for results and checkpoints  
results_folder = Path('./results')  
checkpoint_folder = Path('./checkpoints')  
results_folder.mkdir(exist_ok=True, parents=True)  
checkpoint_folder.mkdir(exist_ok=True, parents=True)  

# Initialize optimizer  
optimizer = AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), eps=1e-8)

# Cosine annealing learning rate scheduler with warmup and minimum LR
num_warmup_steps = 1600 # Number of warmup steps
num_training_steps = 10000  # Total number of training steps (adjust as needed)
min_lr = 1.5e-5  # Minimum learning rate

def get_cosine_schedule_with_warmup_and_min_lr(
    optimizer,
    num_warmup_steps,
    num_training_steps,
    min_lr,
    num_cycles=0.5,
):
    # Get the initial learning rate from the optimizer's parameter group
    base_lr = optimizer.param_groups[0]["lr"]

    def lr_lambda(current_step):
        # Linear warmup phase
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))

        # Cosine annealing phase
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * 2.0 * num_cycles * progress))

        # Scale the cosine decay so that the learning rate never goes below min_lr
        scaled = cosine_decay * (1.0 - min_lr / base_lr) + min_lr / base_lr
        return scaled

    return LambdaLR(optimizer, lr_lambda)

# Instantiate the scheduler using the custom cosine scheduler with warmup and min_lr
scheduler = get_cosine_schedule_with_warmup_and_min_lr(
    optimizer,
    num_warmup_steps=num_warmup_steps,
    num_training_steps=num_training_steps,
    min_lr=min_lr,
    num_cycles=0.5,
)

save_every = 5000
total_steps = 0  

model.train()
progress_bar = tqdm(range(num_training_steps), desc="Training")
for step in progress_bar:

    total_steps += 1  
    progress_bar.set_description(f"Step {total_steps}")
    
    # Gradient accumulation for stability  
    accumulation_steps = 2  
    for _ in range(accumulation_steps):  
        # Get batch of text and images  
        batch = next(iter_dl)  
            
        # Forward pass  
        loss = model(batch)  
            
        # Log loss components  
        if isinstance(loss, dict):  
            for key, value in loss.items():  
                writer.add_scalar(f'Loss/{key}', value.item(), total_steps)  
            loss = loss['total']  # Assuming total loss is returned  
        else:  
            writer.add_scalar('Loss/total', loss.item(), total_steps)  
            
        # Backward pass with scaling  
        (loss / accumulation_steps).backward()  
        
    # Gradient clipping  
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  
        
    # Optimizer step  
    optimizer.step()  
    optimizer.zero_grad()  
    scheduler.step()  
        
    # Update EMA model  
    ema_model.update()  
        
    # Log learning rate  
    writer.add_scalar('LR', scheduler.get_last_lr()[0], total_steps)  
    
    # Regular checkpoint saving  
    if total_steps % save_every == 0:  
        torch.save({  
            'model': model.state_dict(),  
            'ema_model': ema_model.state_dict(),  
            'optimizer': optimizer.state_dict(),  
            'scheduler': scheduler.state_dict(),  
            'step': total_steps,
        }, f"{checkpoint_folder}/checkpoint_{total_steps}.pt")


# === Inference on Test Set ===
if DO_TEST:
    # Load the best model for inference
    test_dataset = COCODataset(
        json_file=TEST_FILE_PATH,
        image_dir=IMAGE_DIR,
        image_size=256
    )
    test_dataloader = model.create_dataloader(test_dataset, batch_size=8, shuffle=False)

    # 创建用于保存测试结果的文件夹
    test_results_folder = Path('./test_results')
    test_results_folder.mkdir(exist_ok=True, parents=True)

    model.eval()
    ema_model.eval()

    print("Starting inference on test set...")
    with torch.no_grad():
        for batch_idx, batch_data in enumerate(tqdm(test_dataloader, desc="Testing")):
            text_tensors = [text_tensor.cuda() for text_tensor, _ in batch_data]
            captions = [decode_tokens(t.cpu()) for t in text_tensors]

            for i, text_tensor in enumerate(text_tensors):
                gen_sample = ema_model.sample([text_tensor])
                print_modality_sample(gen_sample)
                for j, gen_res in enumerate(gen_sample):
                    if isinstance(gen_res, tuple):
                        image = gen_res[1]
                        break

                image = image.detach().cpu().clamp(0, 1)

                # save generated image
                image_path = test_results_folder / f"sample_{batch_idx}_{i}.png"
                save_image(image, image_path)

                # save real caption
                with open(test_results_folder / f"sample_{batch_idx}_{i}.txt", 'w') as f:
                    f.write(captions[i])

    print(f"Inference completed. Results saved to: {test_results_folder}")

[lavinal712]

Transfusion is inherently difficult to train; what you're essentially training is actually a DiT model.