# ------------------------------------------------------------------------------------ # Minimal DALL-E # Copyright (c) 2021 KakaoBrain. All Rights Reserved. # Licensed under the Apache License, Version 2.0 [see LICENSE for details] # ------------------------------------------------------------------------------------ import os import sys import argparse from typing import Optional from datetime import datetime import torch from torch.utils.data import DataLoader import torchvision import torchvision.transforms as transforms import pytorch_lightning as pl from pytorch_lightning.callbacks import ModelCheckpoint, Callback from pytorch_lightning.loggers import TensorBoardLogger from pytorch_lightning.utilities.distributed import rank_zero_only sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from dalle.models import ImageGPT parser = argparse.ArgumentParser() parser.add_argument('-d', '--config-downstream', type=str, default=None, required=True) parser.add_argument('-u', '--path-upstream', type=str, default=None, required=True) parser.add_argument('-r', '--result-path', type=str, default=None, required=True) parser.add_argument('--imagenet-path', type=str, default=None, required=True) parser.add_argument('--n-gpus', type=int, default=1) parser.add_argument('--seed', type=int, default=0) args = parser.parse_args() class ImageLogger(Callback): def __init__(self): super().__init__() @rank_zero_only def log_img(self, pl_module, batch, current_epoch, split="train"): with torch.no_grad(): images, labels = batch recons = pl_module.stage1(images) images = images.cpu() recons = recons.cpu() grid_org = (torchvision.utils.make_grid(images, nrow=8) + 1.0) / 2.0 grid_rec = (torchvision.utils.make_grid(recons, nrow=8) + 1.0) / 2.0 grid_rec = torch.clip(grid_rec, min=0, max=1) pl_module.logger.experiment.add_image(f"images_org/{split}", grid_org, global_step=current_epoch) pl_module.logger.experiment.add_image(f"images_rec/{split}", grid_rec, global_step=current_epoch) def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): if batch_idx == 0 and trainer.current_epoch < 5: self.log_img(pl_module, batch, current_epoch=trainer.current_epoch, split="train") def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): if batch_idx == 0 and trainer.current_epoch < 5: self.log_img(pl_module, batch, current_epoch=trainer.current_epoch, split="test") class ImageNetDataModule(pl.LightningDataModule): def __init__(self, data_dir: Optional[str] = None, image_resolution: int = 256, train_batch_size: int = 2, valid_batch_size: int = 32, num_workers: int = 8): super().__init__() self.data_dir = data_dir self.image_resolution = image_resolution self.train_batch_size = train_batch_size self.valid_batch_size = valid_batch_size self.num_workers = num_workers self.train_transform = transforms.Compose( [transforms.Resize(image_resolution), transforms.RandomCrop(image_resolution), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])] ) self.valid_transform = transforms.Compose( [transforms.Resize(image_resolution), transforms.CenterCrop(image_resolution), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])] ) def setup(self, stage=None): self.trainset = torchvision.datasets.ImageNet(root=self.data_dir, split='train', transform=self.train_transform) self.validset = torchvision.datasets.ImageNet(root=self.data_dir, split='val', transform=self.valid_transform) def train_dataloader(self): return DataLoader(self.trainset, batch_size=self.train_batch_size, num_workers=self.num_workers, pin_memory=True) def valid_dataloader(self): return DataLoader(self.validset, batch_size=self.valid_batch_size, num_workers=self.num_workers, pin_memory=True) def setup_callbacks(config): # Setup callbacks now = datetime.now().strftime('%d%m%Y_%H%M%S') result_path = os.path.join(args.result_path, os.path.basename(args.config_downstream).split('.')[0], now) ckpt_path = os.path.join(result_path, 'ckpt') log_path = os.path.join(result_path, 'log') checkpoint_callback = ModelCheckpoint( dirpath=ckpt_path, filename="imagenet-clscond-gen-{epoch:02d}" if config.stage2.use_cls_cond else "imagenet-uncond-gen-{epoch:02d}", every_n_epochs=config.experiment.save_ckpt_freq, save_weights_only=True, save_last=True ) logger = TensorBoardLogger(log_path, name="iGPT") logger_img = ImageLogger() return checkpoint_callback, logger, logger_img if __name__ == '__main__': pl.seed_everything(args.seed) # Build iGPT model, config = ImageGPT.from_pretrained(args.path_upstream, args.config_downstream) # Setup callbacks ckpt_callback, logger, logger_img = setup_callbacks(config) # Build data modules dataset = ImageNetDataModule(data_dir=args.imagenet_path, image_resolution=config.dataset.image_resolution, train_batch_size=config.experiment.local_batch_size, valid_batch_size=config.experiment.valid_batch_size, num_workers=16) dataset.setup() train_dataloader = dataset.train_dataloader() valid_dataloader = dataset.valid_dataloader() print(f"len(train_dataset) = {len(dataset.trainset)}") print(f"len(valid_dataset) = {len(dataset.validset)}") # Calculate how many batches are accumulated assert config.experiment.total_batch_size % (config.experiment.local_batch_size * args.n_gpus) == 0 grad_accm_steps = config.experiment.total_batch_size // (config.experiment.local_batch_size * args.n_gpus) config.optimizer.max_steps = len(dataset.trainset) // config.experiment.total_batch_size * config.experiment.epochs # Build trainer trainer = pl.Trainer(max_epochs=config.experiment.epochs, accumulate_grad_batches=grad_accm_steps, gradient_clip_val=config.optimizer.grad_clip_norm, precision=16 if config.experiment.use_amp else 32, callbacks=[ckpt_callback, logger_img], accelerator="gpu", devices=args.n_gpus, strategy="ddp", logger=logger) trainer.fit(model, train_dataloader, valid_dataloader)