TwoPerCent/instruct-pix2pix
0
1import argparse, os, sys, datetime, glob2import numpy as np3import time4import torch5import torchvision6import pytorch_lightning as pl7import json8import pickle9 10from packaging import version11from omegaconf import OmegaConf12from torch.utils.data import DataLoader, Dataset13from functools import partial14from PIL import Image15 16import torch.distributed as dist17from pytorch_lightning import seed_everything18from pytorch_lightning.trainer import Trainer19from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor20from pytorch_lightning.utilities.distributed import rank_zero_only21from pytorch_lightning.utilities import rank_zero_info22from pytorch_lightning.plugins import DDPPlugin23 24sys.path.append("./stable_diffusion")25 26from ldm.data.base import Txt2ImgIterableBaseDataset27from ldm.util import instantiate_from_config28 29 30def get_parser(**parser_kwargs):31 def str2bool(v):32 if isinstance(v, bool):33 return v34 if v.lower() in ("yes", "true", "t", "y", "1"):35 return True36 elif v.lower() in ("no", "false", "f", "n", "0"):37 return False38 else:39 raise argparse.ArgumentTypeError("Boolean value expected.")40 41 parser = argparse.ArgumentParser(**parser_kwargs)42 parser.add_argument(43 "-n",44 "--name",45 type=str,46 const=True,47 default="",48 nargs="?",49 help="postfix for logdir",50 )51 parser.add_argument(52 "-r",53 "--resume",54 type=str,55 const=True,56 default="",57 nargs="?",58 help="resume from logdir or checkpoint in logdir",59 )60 parser.add_argument(61 "-b",62 "--base",63 nargs="*",64 metavar="base_config.yaml",65 help="paths to base configs. Loaded from left-to-right. "66 "Parameters can be overwritten or added with command-line options of the form `--key value`.",67 default=list(),68 )69 parser.add_argument(70 "-t",71 "--train",72 type=str2bool,73 const=True,74 default=False,75 nargs="?",76 help="train",77 )78 parser.add_argument(79 "--no-test",80 type=str2bool,81 const=True,82 default=False,83 nargs="?",84 help="disable test",85 )86 parser.add_argument(87 "-p",88 "--project",89 help="name of new or path to existing project"90 )91 parser.add_argument(92 "-d",93 "--debug",94 type=str2bool,95 nargs="?",96 const=True,97 default=False,98 help="enable post-mortem debugging",99 )100 parser.add_argument(101 "-s",102 "--seed",103 type=int,104 default=23,105 help="seed for seed_everything",106 )107 parser.add_argument(108 "-f",109 "--postfix",110 type=str,111 default="",112 help="post-postfix for default name",113 )114 parser.add_argument(115 "-l",116 "--logdir",117 type=str,118 default="logs",119 help="directory for logging dat shit",120 )121 parser.add_argument(122 "--scale_lr",123 action="store_true",124 default=False,125 help="scale base-lr by ngpu * batch_size * n_accumulate",126 )127 return parser128 129 130def nondefault_trainer_args(opt):131 parser = argparse.ArgumentParser()132 parser = Trainer.add_argparse_args(parser)133 args = parser.parse_args([])134 return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k))135 136 137class WrappedDataset(Dataset):138 """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""139 140 def __init__(self, dataset):141 self.data = dataset142 143 def __len__(self):144 return len(self.data)145 146 def __getitem__(self, idx):147 return self.data[idx]148 149 150def worker_init_fn(_):151 worker_info = torch.utils.data.get_worker_info()152 153 dataset = worker_info.dataset154 worker_id = worker_info.id155 156 if isinstance(dataset, Txt2ImgIterableBaseDataset):157 split_size = dataset.num_records // worker_info.num_workers158 # reset num_records to the true number to retain reliable length information159 dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size]160 current_id = np.random.choice(len(np.random.get_state()[1]), 1)161 return np.random.seed(np.random.get_state()[1][current_id] + worker_id)162 else:163 return np.random.seed(np.random.get_state()[1][0] + worker_id)164 165 166class DataModuleFromConfig(pl.LightningDataModule):167 def __init__(self, batch_size, train=None, validation=None, test=None, predict=None,168 wrap=False, num_workers=None, shuffle_test_loader=False, use_worker_init_fn=False,169 shuffle_val_dataloader=False):170 super().__init__()171 self.batch_size = batch_size172 self.dataset_configs = dict()173 self.num_workers = num_workers if num_workers is not None else batch_size * 2174 self.use_worker_init_fn = use_worker_init_fn175 if train is not None:176 self.dataset_configs["train"] = train177 self.train_dataloader = self._train_dataloader178 if validation is not None:179 self.dataset_configs["validation"] = validation180 self.val_dataloader = partial(self._val_dataloader, shuffle=shuffle_val_dataloader)181 if test is not None:182 self.dataset_configs["test"] = test183 self.test_dataloader = partial(self._test_dataloader, shuffle=shuffle_test_loader)184 if predict is not None:185 self.dataset_configs["predict"] = predict186 self.predict_dataloader = self._predict_dataloader187 self.wrap = wrap188 189 def prepare_data(self):190 for data_cfg in self.dataset_configs.values():191 instantiate_from_config(data_cfg)192 193 def setup(self, stage=None):194 self.datasets = dict(195 (k, instantiate_from_config(self.dataset_configs[k]))196 for k in self.dataset_configs)197 if self.wrap:198 for k in self.datasets:199 self.datasets[k] = WrappedDataset(self.datasets[k])200 201 def _train_dataloader(self):202 is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)203 if is_iterable_dataset or self.use_worker_init_fn:204 init_fn = worker_init_fn205 else:206 init_fn = None207 return DataLoader(self.datasets["train"], batch_size=self.batch_size,208 num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True,209 worker_init_fn=init_fn, persistent_workers=True)210 211 def _val_dataloader(self, shuffle=False):212 if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:213 init_fn = worker_init_fn214 else:215 init_fn = None216 return DataLoader(self.datasets["validation"],217 batch_size=self.batch_size,218 num_workers=self.num_workers,219 worker_init_fn=init_fn,220 shuffle=shuffle, persistent_workers=True)221 222 def _test_dataloader(self, shuffle=False):223 is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)224 if is_iterable_dataset or self.use_worker_init_fn:225 init_fn = worker_init_fn226 else:227 init_fn = None228 229 # do not shuffle dataloader for iterable dataset230 shuffle = shuffle and (not is_iterable_dataset)231 232 return DataLoader(self.datasets["test"], batch_size=self.batch_size,233 num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle, persistent_workers=True)234 235 def _predict_dataloader(self, shuffle=False):236 if isinstance(self.datasets['predict'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:237 init_fn = worker_init_fn238 else:239 init_fn = None240 return DataLoader(self.datasets["predict"], batch_size=self.batch_size,241 num_workers=self.num_workers, worker_init_fn=init_fn, persistent_workers=True)242 243 244class SetupCallback(Callback):245 def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config):246 super().__init__()247 self.resume = resume248 self.now = now249 self.logdir = logdir250 self.ckptdir = ckptdir251 self.cfgdir = cfgdir252 self.config = config253 self.lightning_config = lightning_config254 255 def on_keyboard_interrupt(self, trainer, pl_module):256 if trainer.global_rank == 0:257 print("Summoning checkpoint.")258 ckpt_path = os.path.join(self.ckptdir, "last.ckpt")259 trainer.save_checkpoint(ckpt_path)260 261 def on_pretrain_routine_start(self, trainer, pl_module):262 if trainer.global_rank == 0:263 # Create logdirs and save configs264 # os.makedirs(self.logdir, exist_ok=True)265 # os.makedirs(self.ckptdir, exist_ok=True)266 # os.makedirs(self.cfgdir, exist_ok=True)267 268 if "callbacks" in self.lightning_config:269 if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']:270 os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True)271 print("Project config")272 print(OmegaConf.to_yaml(self.config))273 OmegaConf.save(self.config,274 os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)))275 276 print("Lightning config")277 print(OmegaConf.to_yaml(self.lightning_config))278 OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}),279 os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)))280 281def get_world_size():282 if not dist.is_available():283 return 1284 if not dist.is_initialized():285 return 1286 return dist.get_world_size()287 288def all_gather(data):289 """290 Run all_gather on arbitrary picklable data (not necessarily tensors)291 Args:292 data: any picklable object293 Returns:294 list[data]: list of data gathered from each rank295 """296 world_size = get_world_size()297 if world_size == 1:298 return [data]299 300 # serialized to a Tensor301 origin_size = None302 if not isinstance(data, torch.Tensor):303 buffer = pickle.dumps(data)304 storage = torch.ByteStorage.from_buffer(buffer)305 tensor = torch.ByteTensor(storage).to("cuda")306 else:307 origin_size = data.size()308 tensor = data.reshape(-1)309 310 tensor_type = tensor.dtype311 312 # obtain Tensor size of each rank313 local_size = torch.LongTensor([tensor.numel()]).to("cuda")314 size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)]315 dist.all_gather(size_list, local_size)316 size_list = [int(size.item()) for size in size_list]317 max_size = max(size_list)318 319 # receiving Tensor from all ranks320 # we pad the tensor because torch all_gather does not support321 # gathering tensors of different shapes322 tensor_list = []323 for _ in size_list:324 tensor_list.append(torch.FloatTensor(size=(max_size,)).cuda().to(tensor_type))325 if local_size != max_size:326 padding = torch.FloatTensor(size=(max_size - local_size,)).cuda().to(tensor_type)327 tensor = torch.cat((tensor, padding), dim=0)328 dist.all_gather(tensor_list, tensor)329 330 data_list = []331 for size, tensor in zip(size_list, tensor_list):332 if origin_size is None:333 buffer = tensor.cpu().numpy().tobytes()[:size]334 data_list.append(pickle.loads(buffer))335 else:336 buffer = tensor[:size]337 data_list.append(buffer)338 339 if origin_size is not None:340 new_shape = [-1] + list(origin_size[1:])341 resized_list = []342 for data in data_list:343 # suppose the difference of tensor size exist in first dimension344 data = data.reshape(new_shape)345 resized_list.append(data)346 347 return resized_list348 else:349 return data_list350 351class ImageLogger(Callback):352 def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True,353 rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False,354 log_images_kwargs=None):355 super().__init__()356 self.rescale = rescale357 self.batch_freq = batch_frequency358 self.max_images = max_images359 self.logger_log_images = {360 pl.loggers.TestTubeLogger: self._testtube,361 }362 self.log_steps = [2 ** n for n in range(6, int(np.log2(self.batch_freq)) + 1)]363 if not increase_log_steps:364 self.log_steps = [self.batch_freq]365 self.clamp = clamp366 self.disabled = disabled367 self.log_on_batch_idx = log_on_batch_idx368 self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {}369 self.log_first_step = log_first_step370 371 @rank_zero_only372 def _testtube(self, pl_module, images, batch_idx, split):373 for k in images:374 grid = torchvision.utils.make_grid(images[k])375 grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w376 377 tag = f"{split}/{k}"378 pl_module.logger.experiment.add_image(379 tag, grid,380 global_step=pl_module.global_step)381 382 @rank_zero_only383 def log_local(self, save_dir, split, images, prompts,384 global_step, current_epoch, batch_idx):385 root = os.path.join(save_dir, "images", split)386 names = {"reals": "before", "inputs": "after", "reconstruction": "before-vq", "samples": "after-gen"}387 # print(root)388 for k in images:389 grid = torchvision.utils.make_grid(images[k], nrow=8)390 if self.rescale:391 grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w392 grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)393 grid = grid.numpy()394 grid = (grid * 255).astype(np.uint8)395 filename = "gs-{:06}_e-{:06}_b-{:06}_{}.png".format(396 global_step,397 current_epoch,398 batch_idx,399 names[k])400 path = os.path.join(root, filename)401 os.makedirs(os.path.split(path)[0], exist_ok=True)402 # print(path)403 Image.fromarray(grid).save(path)404 405 filename = "gs-{:06}_e-{:06}_b-{:06}_prompt.json".format(406 global_step,407 current_epoch,408 batch_idx)409 path = os.path.join(root, filename)410 with open(path, "w") as f:411 for p in prompts:412 f.write(f"{json.dumps(p)}\n")413 414 def log_img(self, pl_module, batch, batch_idx, split="train"):415 check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step416 if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0417 hasattr(pl_module, "log_images") and418 callable(pl_module.log_images) and419 self.max_images > 0) or (split == "val" and batch_idx == 0):420 logger = type(pl_module.logger)421 422 is_train = pl_module.training423 if is_train:424 pl_module.eval()425 426 with torch.no_grad():427 images = pl_module.log_images(batch, split=split, **self.log_images_kwargs)428 429 prompts = batch["edit"]["c_crossattn"][:self.max_images]430 prompts = [p for ps in all_gather(prompts) for p in ps]431 432 for k in images:433 N = min(images[k].shape[0], self.max_images)434 images[k] = images[k][:N]435 images[k] = torch.cat(all_gather(images[k][:N]))436 if isinstance(images[k], torch.Tensor):437 images[k] = images[k].detach().cpu()438 if self.clamp:439 images[k] = torch.clamp(images[k], -1., 1.)440 441 self.log_local(pl_module.logger.save_dir, split, images, prompts,442 pl_module.global_step, pl_module.current_epoch, batch_idx)443 444 logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None)445 logger_log_images(pl_module, images, pl_module.global_step, split)446 447 if is_train:448 pl_module.train()449 450 def check_frequency(self, check_idx):451 if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (452 check_idx > 0 or self.log_first_step):453 if len(self.log_steps) > 0:454 self.log_steps.pop(0)455 return True456 return False457 458 def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):459 if not self.disabled and (pl_module.global_step > 0 or self.log_first_step):460 self.log_img(pl_module, batch, batch_idx, split="train")461 462 def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):463 if not self.disabled and pl_module.global_step > 0:464 self.log_img(pl_module, batch, batch_idx, split="val")465 if hasattr(pl_module, 'calibrate_grad_norm'):466 if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0:467 self.log_gradients(trainer, pl_module, batch_idx=batch_idx)468 469 470class CUDACallback(Callback):471 # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py472 def on_train_epoch_start(self, trainer, pl_module):473 # Reset the memory use counter474 torch.cuda.reset_peak_memory_stats(trainer.root_gpu)475 torch.cuda.synchronize(trainer.root_gpu)476 self.start_time = time.time()477 478 def on_train_epoch_end(self, trainer, pl_module, outputs):479 torch.cuda.synchronize(trainer.root_gpu)480 max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20481 epoch_time = time.time() - self.start_time482 483 try:484 max_memory = trainer.training_type_plugin.reduce(max_memory)485 epoch_time = trainer.training_type_plugin.reduce(epoch_time)486 487 rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds")488 rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB")489 except AttributeError:490 pass491 492 493if __name__ == "__main__":494 # custom parser to specify config files, train, test and debug mode,495 # postfix, resume.496 # `--key value` arguments are interpreted as arguments to the trainer.497 # `nested.key=value` arguments are interpreted as config parameters.498 # configs are merged from left-to-right followed by command line parameters.499 500 # model:501 # base_learning_rate: float502 # target: path to lightning module503 # params:504 # key: value505 # data:506 # target: main.DataModuleFromConfig507 # params:508 # batch_size: int509 # wrap: bool510 # train:511 # target: path to train dataset512 # params:513 # key: value514 # validation:515 # target: path to validation dataset516 # params:517 # key: value518 # test:519 # target: path to test dataset520 # params:521 # key: value522 # lightning: (optional, has sane defaults and can be specified on cmdline)523 # trainer:524 # additional arguments to trainer525 # logger:526 # logger to instantiate527 # modelcheckpoint:528 # modelcheckpoint to instantiate529 # callbacks:530 # callback1:531 # target: importpath532 # params:533 # key: value534 535 now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")536 537 # add cwd for convenience and to make classes in this file available when538 # running as `python main.py`539 # (in particular `main.DataModuleFromConfig`)540 sys.path.append(os.getcwd())541 542 parser = get_parser()543 parser = Trainer.add_argparse_args(parser)544 545 opt, unknown = parser.parse_known_args()546 547 assert opt.name548 cfg_fname = os.path.split(opt.base[0])[-1]549 cfg_name = os.path.splitext(cfg_fname)[0]550 nowname = f"{cfg_name}_{opt.name}"551 logdir = os.path.join(opt.logdir, nowname)552 ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")553 resume = False554 555 if os.path.isfile(ckpt):556 opt.resume_from_checkpoint = ckpt557 base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml")))558 opt.base = base_configs + opt.base559 _tmp = logdir.split("/")560 nowname = _tmp[-1]561 resume = True562 563 ckptdir = os.path.join(logdir, "checkpoints")564 cfgdir = os.path.join(logdir, "configs")565 566 os.makedirs(logdir, exist_ok=True)567 os.makedirs(ckptdir, exist_ok=True)568 os.makedirs(cfgdir, exist_ok=True)569 570 try:571 # init and save configs572 configs = [OmegaConf.load(cfg) for cfg in opt.base]573 cli = OmegaConf.from_dotlist(unknown)574 config = OmegaConf.merge(*configs, cli)575 576 if resume:577 # By default, when finetuning from Stable Diffusion, we load the EMA-only checkpoint to initialize all weights.578 # If resuming InstructPix2Pix from a finetuning checkpoint, instead load both EMA and non-EMA weights.579 config.model.params.load_ema = True580 581 lightning_config = config.pop("lightning", OmegaConf.create())582 # merge trainer cli with config583 trainer_config = lightning_config.get("trainer", OmegaConf.create())584 # default to ddp585 trainer_config["accelerator"] = "ddp"586 for k in nondefault_trainer_args(opt):587 trainer_config[k] = getattr(opt, k)588 if not "gpus" in trainer_config:589 del trainer_config["accelerator"]590 cpu = True591 else:592 gpuinfo = trainer_config["gpus"]593 print(f"Running on GPUs {gpuinfo}")594 cpu = False595 trainer_opt = argparse.Namespace(**trainer_config)596 lightning_config.trainer = trainer_config597 598 # model599 model = instantiate_from_config(config.model)600 601 # trainer and callbacks602 trainer_kwargs = dict()603 604 # default logger configs605 default_logger_cfgs = {606 "wandb": {607 "target": "pytorch_lightning.loggers.WandbLogger",608 "params": {609 "name": nowname,610 "save_dir": logdir,611 "id": nowname,612 }613 },614 "testtube": {615 "target": "pytorch_lightning.loggers.TestTubeLogger",616 "params": {617 "name": "testtube",618 "save_dir": logdir,619 }620 },621 }622 default_logger_cfg = default_logger_cfgs["wandb"]623 if "logger" in lightning_config:624 logger_cfg = lightning_config.logger625 else:626 logger_cfg = OmegaConf.create()627 logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg)628 trainer_kwargs["logger"] = instantiate_from_config(logger_cfg)629 630 # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to631 # specify which metric is used to determine best models632 default_modelckpt_cfg = {633 "target": "pytorch_lightning.callbacks.ModelCheckpoint",634 "params": {635 "dirpath": ckptdir,636 "filename": "{epoch:06}",637 "verbose": True,638 "save_last": True,639 }640 }641 642 if "modelcheckpoint" in lightning_config:643 modelckpt_cfg = lightning_config.modelcheckpoint644 else:645 modelckpt_cfg = OmegaConf.create()646 modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg)647 print(f"Merged modelckpt-cfg: \n{modelckpt_cfg}")648 if version.parse(pl.__version__) < version.parse('1.4.0'):649 trainer_kwargs["checkpoint_callback"] = instantiate_from_config(modelckpt_cfg)650 651 # add callback which sets up log directory652 default_callbacks_cfg = {653 "setup_callback": {654 "target": "main.SetupCallback",655 "params": {656 "resume": opt.resume,657 "now": now,658 "logdir": logdir,659 "ckptdir": ckptdir,660 "cfgdir": cfgdir,661 "config": config,662 "lightning_config": lightning_config,663 }664 },665 "image_logger": {666 "target": "main.ImageLogger",667 "params": {668 "batch_frequency": 750,669 "max_images": 4,670 "clamp": True671 }672 },673 "learning_rate_logger": {674 "target": "main.LearningRateMonitor",675 "params": {676 "logging_interval": "step",677 # "log_momentum": True678 }679 },680 "cuda_callback": {681 "target": "main.CUDACallback"682 },683 }684 if version.parse(pl.__version__) >= version.parse('1.4.0'):685 default_callbacks_cfg.update({'checkpoint_callback': modelckpt_cfg})686 687 if "callbacks" in lightning_config:688 callbacks_cfg = lightning_config.callbacks689 else:690 callbacks_cfg = OmegaConf.create()691 692 print(693 'Caution: Saving checkpoints every n train steps without deleting. This might require some free space.')694 default_metrics_over_trainsteps_ckpt_dict = {695 'metrics_over_trainsteps_checkpoint': {696 "target": 'pytorch_lightning.callbacks.ModelCheckpoint',697 'params': {698 "dirpath": os.path.join(ckptdir, 'trainstep_checkpoints'),699 "filename": "{epoch:06}-{step:09}",700 "verbose": True,701 'save_top_k': -1,702 'every_n_train_steps': 1000,703 'save_weights_only': True704 }705 }706 }707 default_callbacks_cfg.update(default_metrics_over_trainsteps_ckpt_dict)708 709 callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg)710 if 'ignore_keys_callback' in callbacks_cfg and hasattr(trainer_opt, 'resume_from_checkpoint'):711 callbacks_cfg.ignore_keys_callback.params['ckpt_path'] = trainer_opt.resume_from_checkpoint712 elif 'ignore_keys_callback' in callbacks_cfg:713 del callbacks_cfg['ignore_keys_callback']714 715 trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg]716 717 trainer = Trainer.from_argparse_args(trainer_opt, plugins=DDPPlugin(find_unused_parameters=False), **trainer_kwargs)718 trainer.logdir = logdir ###719 720 # data721 data = instantiate_from_config(config.data)722 # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html723 # calling these ourselves should not be necessary but it is.724 # lightning still takes care of proper multiprocessing though725 data.prepare_data()726 data.setup()727 print("#### Data #####")728 for k in data.datasets:729 print(f"{k}, {data.datasets[k].__class__.__name__}, {len(data.datasets[k])}")730 731 # configure learning rate732 bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate733 if not cpu:734 ngpu = len(lightning_config.trainer.gpus.strip(",").split(','))735 else:736 ngpu = 1737 if 'accumulate_grad_batches' in lightning_config.trainer:738 accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches739 else:740 accumulate_grad_batches = 1741 print(f"accumulate_grad_batches = {accumulate_grad_batches}")742 lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches743 if opt.scale_lr:744 model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr745 print(746 "Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (base_lr)".format(747 model.learning_rate, accumulate_grad_batches, ngpu, bs, base_lr))748 else:749 model.learning_rate = base_lr750 print("++++ NOT USING LR SCALING ++++")751 print(f"Setting learning rate to {model.learning_rate:.2e}")752 753 754 # allow checkpointing via USR1755 def melk(*args, **kwargs):756 # run all checkpoint hooks757 if trainer.global_rank == 0:758 print("Summoning checkpoint.")759 ckpt_path = os.path.join(ckptdir, "last.ckpt")760 trainer.save_checkpoint(ckpt_path)761 762 763 def divein(*args, **kwargs):764 if trainer.global_rank == 0:765 import pudb;766 pudb.set_trace()767 768 769 import signal770 771 signal.signal(signal.SIGUSR1, melk)772 signal.signal(signal.SIGUSR2, divein)773 774 # run775 if opt.train:776 try:777 trainer.fit(model, data)778 except Exception:779 melk()780 raise781 if not opt.no_test and not trainer.interrupted:782 trainer.test(model, data)783 except Exception:784 if opt.debug and trainer.global_rank == 0:785 try:786 import pudb as debugger787 except ImportError:788 import pdb as debugger789 debugger.post_mortem()790 raise791 finally:792 # move newly created debug project to debug_runs793 if opt.debug and not opt.resume and trainer.global_rank == 0:794 dst, name = os.path.split(logdir)795 dst = os.path.join(dst, "debug_runs", name)796 os.makedirs(os.path.split(dst)[0], exist_ok=True)797 os.rename(logdir, dst)798 if trainer.global_rank == 0:799 print(trainer.profiler.summary())800