train_vqgan.py

import os
import time

import hydra
import jittor as jt
import jittor.nn as nn

from omegaconf import OmegaConf

import wandb
from model_jittor.dataset import get_vq_dataloader
from model_jittor.autoencoder.vqgan import VQModel
from model_jittor.autoencoder.loss import (
    NLayerDiscriminator, 
    VQAECriterion, DiscriminatorCriterion, BaseCriterion)
from utils import save_checkpoint, log, start_grad, stop_grad

from typing import Optional


DEBUG = False

@hydra.main(version_base=None, config_path='configs', config_name='vqgan_classic.yaml')
def main(cfg):
    global best_rec_loss

    # initialization
    config = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True)
    
    jt.flags.use_cuda=True
    jt.flags.auto_mixed_precision_level = 0
    jt.set_global_seed(cfg.seed)
    
    cfg.data.batch_size *= jt.world_size
    # TODO: adjust lr with respect bs.
    print(f"Using {jt.world_size} gpu(s), current batch size {cfg.data.batch_size}")
    
    # TODO: add resume for wandb
    if jt.rank == 0:
        wandb.init(project=cfg.project, name=cfg.name, config=config)
        # save ckpt in './save/run-id/checkpoints', run-id is generated by wandb
        cfg.save_dir = wandb.run.dir.replace('wandb', 'save'
                                            ).replace( 'files', 'checkpoints')
        os.makedirs(os.path.join(cfg.save_dir), exist_ok=True)
        print(f'Saving checkpoints in {cfg.save_dir}')

    # data
    train_dataloader, val_dataloader = get_vq_dataloader(**cfg.data)

    # create model and summary
    model = VQModel(**cfg.model)
    discriminator = NLayerDiscriminator(**cfg.disc)

    # configure optimizer TODO: add lr_scheduler
    optimizer_ae = jt.optim.Adam(
        model.parameters(), lr=cfg.lr, betas=(0.5, 0.9),
    )
    
    optimizer_disc= jt.optim.Adam(
        discriminator.parameters(), lr=cfg.lr, betas=(0.5, 0.9),
    )

    # resume 
    if cfg.resume is not None:
        assert os.path.isfile(cfg.resume)
        checkpoint = jt.load(cfg.resume)
        model.load_state_dict(checkpoint['model'])
        if 'discriminator' in checkpoint and 'optimizer_disc' in checkpoint:
            optimizer_disc.load_state_dict(checkpoint['optimizer_disc'])
            discriminator.load_state_dict(checkpoint['discriminator'])
        optimizer_ae.load_state_dict(checkpoint['optimizer_ae'])
        cfg.start_epoch = checkpoint['epoch'] + 1  # start from the next epoch
        # NOTE: best_rec_loss may have to be move .to(device)
        best_rec_loss = checkpoint['best_rec_loss']
    
    best_rec_loss = 1e4
    
    gen_criterion = VQAECriterion(**cfg.loss)
    disc_criterion = DiscriminatorCriterion(**cfg.loss)

    print('Start training, good luck!')
    for epoch in range(cfg.start_epoch, cfg.epochs):

        start_time = time.time()
        train_one_epoch(model=model,
                        discriminator=discriminator,
                        train_dataset=train_dataloader, 
                        optimizer_ae=optimizer_ae, 
                        optimizer_disc=optimizer_disc,
                        epoch=epoch, 
                        cfg=cfg,
                        gen_criterion=gen_criterion,
                        disc_criterion=disc_criterion,
                        with_gan_loss=epoch > cfg.loss.gan_start)
        train_time = time.time() - start_time
        if jt.rank == 0:
            print(f'Epoch {epoch:3d} training time: {train_time/60:.2f} min.')
        
        rec_loss = validate(model=model, 
                            discrimiator=discriminator,
                            val_dataset=val_dataloader, 
                            epoch=epoch, 
                            cfg=cfg,
                            gen_criterion=gen_criterion,
                            disc_criterion=disc_criterion)
        total_time = time.time() - start_time
        if jt.rank == 0:
            print(f'Epoch {epoch:3d} total time: {total_time/60:.2f} min.')
        
        if jt.rank == 0:
            # save
            if rec_loss < best_rec_loss:
                best_rec_loss = rec_loss
                is_best=True
            else: 
                is_best=False
            
            if is_best or epoch % cfg.save_freq == 0:
                save_checkpoint({  
                        'epoch': epoch,
                        'model': model.state_dict(),
                        'discriminator': discriminator.state_dict(),
                        'optimizer_ae': optimizer_ae.state_dict(),
                        'optimizer_disc': optimizer_disc.state_dict(),
                        'best_rec_loss': best_rec_loss,
                    },
                    is_best=is_best,
                    save_dir=cfg.save_dir,
                    filename=f"epoch_{epoch}.ckpt",
                )

def train_one_epoch(model: VQModel, 
                    optimizer_ae: jt.optim.Optimizer, 
                    train_dataset: jt.dataset.Dataset, 
                    epoch: int, 
                    cfg,
                    discriminator: nn.Module = None,
                    optimizer_disc: jt.optim.Optimizer = None,
                    gen_criterion: VQAECriterion = None,
                    disc_criterion: DiscriminatorCriterion = None,
                    with_gan_loss=False
                ): # TODO: support criterion
    model.train()
    discriminator.train()
    for i, (image, name) in enumerate(train_dataset):
        if DEBUG and i == 50: break
        global_train_steps = epoch * len(train_dataset) + i
        all_logs = {}
        
        image_rec, q_loss, _ = model(image)
        
        # train discriminator
        start_grad(discriminator)
        loss, logs = disc_criterion(image_rec, image, model, discriminator,
                                    with_gan_loss=with_gan_loss)
        optimizer_disc.step(loss)
        all_logs.update(logs)
        
        # train model
        stop_grad(discriminator)
        loss, logs = gen_criterion(image_rec, image, model, discriminator, 
                                   q_loss=q_loss,
                                   with_gan_loss=with_gan_loss)
        optimizer_ae.step(loss)
        all_logs.update(logs)
        
        jt.sync_all()
        
        log(all_logs, epoch, i, global_train_steps, len(train_dataset),
            images=dict(image=image, image_rec=image_rec),
            log_interval=cfg.print_freq,
            image_interval=cfg.save_wandb_image_freq)
 
 
def validate(model: VQModel, 
            val_dataset: jt.dataset.Dataset, 
            epoch: int, 
            cfg,
            discrimiator: Optional[nn.Module] = None,
            gen_criterion: Optional[BaseCriterion] = None,
            disc_criterion: Optional[BaseCriterion]  = None
        ): # TODO: support criterion
    rec_losses = 0
    model.eval()
    discrimiator.eval()

    with jt.no_grad():
        for i, (image, name) in enumerate(val_dataset):
            if DEBUG and i == 50: break
            
            global_val_steps = epoch * len(val_dataset) + i
            
            image = image.stop_grad()
            image_rec, q_loss, _ = model(image)
            
            all_logs = {}

            loss, logs = disc_criterion(image_rec, image, model, discrimiator,
                                        with_gan_loss=True)
            all_logs.update(logs)
            
            loss, logs = gen_criterion(image_rec, image, model, discrimiator,
                                       q_loss=q_loss,
                                       with_gan_loss=True)
            all_logs.update(logs)
            
            rec_losses += all_logs["rec_loss"]

            jt.sync_all(True)
            
            log(all_logs, epoch, i, global_val_steps, len(val_dataset), stage="val",
                images=dict(image=image, image_rec=image_rec),
                log_interval=cfg.print_freq,
                image_interval=cfg.save_wandb_image_freq)
    
    jt.sync_all()
    if jt.in_mpi:
        rec_losses = rec_losses.mpi_all_reduce()
    rec_avg_loss = rec_losses.data[0] / val_dataset.total_len
    return rec_avg_loss


if __name__ == "__main__":
    main()