import argparse, os, sys, datetime, glob os.environ["TOKENIZERS_PARALLELISM"] = "false" import numpy as np import time import torch import torch.distributed as dist import torchvision import pytorch_lightning as pl import matplotlib.pyplot as plt import soundfile from omegaconf import OmegaConf from torch.utils.data import DataLoader, Dataset from functools import partial import ldm from pytorch_lightning import seed_everything from pytorch_lightning.trainer import Trainer from pytorch_lightning.callbacks import ModelCheckpoint, Callback,LearningRateMonitor from pytorch_lightning.utilities.distributed import rank_zero_only from pytorch_lightning.utilities import rank_zero_info from ldm.util import instantiate_from_config from ldm.data.joinaudiodataset_anylen import JoinManifestSpecs from ldm.data.joinaudiodataset_struct_sample_anylen import JoinManifestSpecs def get_parser(**parser_kwargs): def str2bool(v): if isinstance(v, bool): return v if v.lower() in ("yes", "true", "t", "y", "1"): return True elif v.lower() in ("no", "false", "f", "n", "0"): return False else: raise argparse.ArgumentTypeError("Boolean value expected.") parser = argparse.ArgumentParser(**parser_kwargs) parser.add_argument( "-n", "--name", type=str, const=True, default="", nargs="?", help="postfix for logdir", ) parser.add_argument( "-r", "--resume", type=str, const=True, default="", nargs="?", help="resume from logdir or checkpoint in logdir", ) parser.add_argument( "-b", "--base", nargs="*", metavar="base_config.yaml", help="paths to base configs. Loaded from left-to-right. " "Parameters can be overwritten or added with command-line options of the form `--key value`.", default=list(), ) parser.add_argument( "-t", "--train", type=str2bool, const=True, default=False, nargs="?", help="train", ) parser.add_argument( "-val", type=str2bool, const=True, default=False, nargs="?", help="validation", ) parser.add_argument( "--no-test", type=str2bool, const=True, default=False, nargs="?", help="disable test", ) parser.add_argument( "--test-repeat", type=int, default=1, help="repeat each caption for t times in test", ) parser.add_argument( "-p", "--project", help="name of new or path to existing project" ) parser.add_argument( "-d", "--debug", type=str2bool, nargs="?", const=True, default=False, help="enable post-mortem debugging", ) parser.add_argument( "-s", "--seed", type=int, default=23, help="seed for seed_everything", ) parser.add_argument( "-f", "--postfix", type=str, default="", help="post-postfix for default name", ) parser.add_argument( "-l", "--logdir", type=str, default="logs", help="directory for logging dat shit", ) parser.add_argument( "--scale_lr", type=str2bool, nargs="?", const=True, default=True, help="scale base-lr by ngpu * batch_size * n_accumulate", ) return parser def getrank(): def is_dist_avail_and_initialized(): if not dist.is_available(): return False if not dist.is_initialized(): return False return True if not is_dist_avail_and_initialized(): return 0 return dist.get_rank() def nondefault_trainer_args(opt): parser = argparse.ArgumentParser() parser = Trainer.add_argparse_args(parser) args = parser.parse_args([]) return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k)) class WrappedDataset(Dataset): """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset""" def __init__(self, dataset): self.data = dataset def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] def worker_init_fn(_): worker_info = torch.utils.data.get_worker_info() dataset = worker_info.dataset worker_id = worker_info.id return np.random.seed(np.random.get_state()[1][0] + worker_id) class DataModuleFromConfig(pl.LightningDataModule):# batchloader outputshape should be (b,h,w,c) and it will be permuted to (b,c,h,w) in autoencoder.get_input() def __init__(self, batch_size, train=None, validation=None, test=None, predict=None, wrap=False, num_workers=None, shuffle_test_loader=False, use_worker_init_fn=False, shuffle_val_dataloader=False): super().__init__() self.batch_size = batch_size self.dataset_configs = dict() self.num_workers = num_workers if num_workers is not None else batch_size * 2 self.use_worker_init_fn = use_worker_init_fn if train is not None: self.dataset_configs["train"] = train self.train_dataloader = self._train_dataloader if validation is not None: self.dataset_configs["validation"] = validation self.val_dataloader = partial(self._val_dataloader, shuffle=shuffle_val_dataloader) if test is not None: self.dataset_configs["test"] = test self.test_dataloader = partial(self._test_dataloader, shuffle=shuffle_test_loader) if predict is not None: self.dataset_configs["predict"] = predict self.predict_dataloader = self._predict_dataloader self.wrap = wrap def prepare_data(self): for data_cfg in self.dataset_configs.values(): instantiate_from_config(data_cfg) def setup(self, stage=None): self.datasets = dict( (k, instantiate_from_config(self.dataset_configs[k])) for k in self.dataset_configs) if self.wrap: for k in self.datasets: self.datasets[k] = WrappedDataset(self.datasets[k]) def _train_dataloader(self): init_fn = None if isinstance(self.datasets["train"],ldm.data.joinaudiodataset_anylen.JoinManifestSpecs): from ldm.data.joinaudiodataset_anylen import DDPIndexBatchSampler dataset = self.datasets["train"] batch_sampler = DDPIndexBatchSampler(indices=dataset.ordered_indices(),batch_size=self.batch_size,shuffle=True,drop_last=True) return DataLoader(dataset, batch_sampler=batch_sampler,sampler=None, num_workers=self.num_workers, collate_fn=dataset.collater, worker_init_fn=init_fn) elif isinstance(self.datasets["train"],ldm.data.joinaudiodataset_struct_anylen.JoinManifestSpecs): from ldm.data.joinaudiodataset_struct_anylen import DDPIndexBatchSampler dataset = self.datasets["train"] batch_sampler = DDPIndexBatchSampler(indices=dataset.ordered_indices(),batch_size=self.batch_size,shuffle=True,drop_last=True) return DataLoader(dataset, batch_sampler=batch_sampler,sampler=None, num_workers=self.num_workers, collate_fn=dataset.collater, worker_init_fn=init_fn) elif isinstance(self.datasets["train"],ldm.data.joinaudiodataset_struct_sample_anylen.JoinManifestSpecs): from ldm.data.joinaudiodataset_struct_sample_anylen import DDPIndexBatchSampler dataset = self.datasets["train"] main_indices,other_indices = dataset.ordered_indices() # main_indices = dataset.ordered_indices() batch_sampler = DDPIndexBatchSampler(main_indices,other_indices,batch_size=self.batch_size,shuffle=True,drop_last=True) # batch_sampler = DDPIndexBatchSampler(main_indices,batch_size=self.batch_size,shuffle=True,drop_last=True) loader = DataLoader(dataset, batch_sampler=batch_sampler,sampler=None, num_workers=self.num_workers, collate_fn=dataset.collater, worker_init_fn=init_fn) print("train_loader_length",len(loader)) return loader else: return DataLoader(self.datasets["train"], batch_size=self.batch_size ,# sampler=DistributedSampler # np.arange(100), num_workers=self.num_workers, shuffle=True, worker_init_fn=init_fn) def _val_dataloader(self, shuffle=False): init_fn = None if isinstance(self.datasets["validation"],ldm.data.joinaudiodataset_struct_anylen.JoinManifestSpecs): from ldm.data.joinaudiodataset_struct_anylen import DDPIndexBatchSampler dataset = self.datasets["validation"] batch_sampler = DDPIndexBatchSampler(indices=dataset.ordered_indices(),batch_size=self.batch_size,shuffle=shuffle,drop_last=True) return DataLoader(dataset, batch_sampler=batch_sampler,sampler=None, num_workers=self.num_workers, collate_fn=dataset.collater, worker_init_fn=init_fn) if isinstance(self.datasets["validation"],JoinManifestSpecs): from ldm.data.joinaudiodataset_struct_sample_anylen import DDPIndexBatchSampler dataset = self.datasets["validation"] main_indices,other_indices = dataset.ordered_indices() batch_sampler = DDPIndexBatchSampler(main_indices,other_indices,batch_size=self.batch_size,shuffle=shuffle,drop_last=True) return DataLoader(dataset, batch_sampler=batch_sampler,sampler=None, num_workers=self.num_workers, collate_fn=dataset.collater, worker_init_fn=init_fn) else: return DataLoader(self.datasets["validation"], batch_size=self.batch_size, num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle) def _test_dataloader(self, shuffle=False): init_fn = None # do not shuffle dataloader for iterable dataset return DataLoader(self.datasets["test"], batch_size=self.batch_size, num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle) def _predict_dataloader(self, shuffle=False): init_fn = None return DataLoader(self.datasets["predict"], batch_size=self.batch_size, num_workers=self.num_workers, worker_init_fn=init_fn) class SpectrogramDataModuleFromConfig(DataModuleFromConfig): '''avoiding duplication of hyper-parameters in the config by gross patching here ''' def __init__(self, batch_size, num_workers,spec_dir_path=None,main_spec_dir_path=None,other_spec_dir_path=None, mel_num=None, spec_len=None, spec_crop_len=1248,drop=0,mode='pad', require_caption=True, train=None, validation=None, test=None, predict=None, wrap=False): specs_dataset_cfg = { 'spec_dir_path': spec_dir_path, 'main_spec_dir_path':main_spec_dir_path, 'other_spec_dir_path':other_spec_dir_path, 'require_caption': require_caption, 'mel_num': mel_num, 'spec_len': spec_len, 'spec_crop_len': spec_crop_len, 'mode': mode, 'drop': drop } for name, split in {'train': train, 'validation': validation, 'test': test}.items(): if split is not None: split.params.specs_dataset_cfg = specs_dataset_cfg super().__init__(batch_size, train, validation, test, predict, wrap, num_workers) class SetupCallback(Callback):# will not load ckpt, just set directories for the experiment def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config): super().__init__() self.resume = resume self.now = now self.logdir = logdir self.ckptdir = ckptdir self.cfgdir = cfgdir self.config = config self.lightning_config = lightning_config def on_exception(self, trainer, pl_module, exception): if trainer.global_rank == 0: print("Summoning checkpoint.") ckpt_path = os.path.join(self.ckptdir, "last.ckpt") trainer.save_checkpoint(ckpt_path) def on_fit_start(self, trainer, pl_module): if trainer.global_rank == 0: # Create logdirs and save configs os.makedirs(self.logdir, exist_ok=True) os.makedirs(self.ckptdir, exist_ok=True) os.makedirs(self.cfgdir, exist_ok=True) if "callbacks" in self.lightning_config: if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']: os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True) print("Project config") print(OmegaConf.to_yaml(self.config)) OmegaConf.save(self.config, os.path.join(self.cfgdir, "{}-project.yaml".format(self.now))) print("Lightning config") print(OmegaConf.to_yaml(self.lightning_config)) OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}), os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now))) else: # ModelCheckpoint callback created log directory --- remove it if not self.resume and os.path.exists(self.logdir): dst, name = os.path.split(self.logdir) dst = os.path.join(dst, "child_runs", name) os.makedirs(os.path.split(dst)[0], exist_ok=True) try: os.rename(self.logdir, dst) except FileNotFoundError: pass class ImageLogger(Callback): def __init__(self, batch_frequency, max_images, increase_log_steps=True, disabled=False, log_on_batch_idx=False, log_first_step=False,melvmin=0,melvmax=1, log_images_kwargs=None,**kwargs): super().__init__() self.batch_freq = batch_frequency self.max_images = max_images self.logger_log_images = { pl.loggers.TensorBoardLogger: self._log, } self.log_steps = [2 ** n for n in range(int(np.log2(self.batch_freq)) + 1)] if not increase_log_steps: self.log_steps = [self.batch_freq] self.disabled = disabled self.log_on_batch_idx = log_on_batch_idx self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} self.log_first_step = log_first_step self.melvmin=melvmin self.melvmax=melvmax @rank_zero_only def _log(self, pl_module, images, batch_idx, split): for k in images: grid = torchvision.utils.make_grid(images[k]) fig = plt.figure() plt.pcolor(grid.mean(dim=0),vmin=self.melvmin,vmax=self.melvmax) tag = f"{split}/{k}" pl_module.logger.experiment.add_figure(tag, fig,global_step=pl_module.global_step) @rank_zero_only def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): root = os.path.join(save_dir, "images", split) for k in images: grid = torchvision.utils.make_grid(images[k], nrow=4)# c=3,h,w grid = grid.mean(dim=0)# to 1 channel grid = grid.numpy() filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( k, global_step, current_epoch, batch_idx) path = os.path.join(root, filename) os.makedirs(os.path.split(path)[0], exist_ok=True) plt.imsave(path,grid,vmin=self.melvmin,vmax=self.melvmax) def log_img(self, pl_module, batch, batch_idx, split="train"): check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0 hasattr(pl_module, "log_images") and callable(pl_module.log_images) and self.max_images > 0): logger = type(pl_module.logger) is_train = pl_module.training if is_train: pl_module.eval() with torch.no_grad():# 这里会调用ddpm中的log_images images = pl_module.log_images(batch, split=split, **self.log_images_kwargs)# images is a dict for k in images.keys(): N = min(images[k].shape[0], self.max_images) images[k] = images[k][:N] if isinstance(images[k], torch.Tensor): images[k] = images[k].detach().cpu() self.log_local(pl_module.logger.save_dir, split, images, pl_module.global_step, pl_module.current_epoch, batch_idx) logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None) logger_log_images(pl_module, images, pl_module.global_step, split) if is_train: pl_module.train() def check_frequency(self, check_idx): if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and ( check_idx > 0 or self.log_first_step): return True return False def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): if not self.disabled and (pl_module.global_step > 0 or self.log_first_step): self.log_img(pl_module, batch, batch_idx, split="train") # pass def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx,dataloader_idx): if not self.disabled and pl_module.global_step > 0: self.log_img(pl_module, batch, batch_idx, split="val") if hasattr(pl_module, 'calibrate_grad_norm'): if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0: self.log_gradients(trainer, pl_module, batch_idx=batch_idx) class AudioLogger(ImageLogger): def __init__(self, batch_frequency, max_images, increase_log_steps=True, melvmin=0,melvmax=1,disabled=False, log_on_batch_idx=False, log_first_step=False, log_images_kwargs=None, for_specs=False, vocoder_cfg=None, spec_dir_name=None, sample_rate=None,**kwargs): super().__init__(batch_frequency, max_images, increase_log_steps, disabled, log_on_batch_idx, log_first_step, melvmin,melvmax,log_images_kwargs) self.for_specs = for_specs self.spec_dir_name = spec_dir_name self.sample_rate = sample_rate print('We will not save audio for conditioning and conditioning_rec') if self.for_specs: self.vocoder = instantiate_from_config(vocoder_cfg) def _visualize_attention(self, attention, scale_by_prior=True): if scale_by_prior: B, H, T, T = attention.shape # attention weight is 1/T: if we have a seq with length 3 the weights are 1/3, 1/3, and 1/3 # making T by T matrix with zeros in the upper triangular part attention_uniform_prior = 1 / torch.arange(1, T+1).view(1, T, 1).repeat(B, 1, T) attention_uniform_prior = attention_uniform_prior.tril().view(B, 1, T, T).to(attention.device) attention = attention - attention_uniform_prior attention_agg = attention.sum(dim=1, keepdims=True) return attention_agg def _log_rec_audio(self, specs, tag, global_step, pl_module=None, save_rec_path=None): # specs are (B, 1, F, T) for i, spec in enumerate(specs): spec = spec.data.squeeze(0).cpu().numpy() if spec.shape[0] != 80: continue wave = self.vocoder.vocode(spec) wave = torch.from_numpy(wave).unsqueeze(0) if pl_module is not None: pl_module.logger.experiment.add_audio(f'{tag}_{i}', wave, global_step, self.sample_rate) # in case we would like to save it on disk if save_rec_path is not None: soundfile.write(save_rec_path, wave.squeeze(0).numpy(), self.sample_rate, 'FLOAT') @rank_zero_only def _log(self, pl_module, images, batch_idx, split): for k in images: # images is a dict,images[k]'s shape is (B,C,H,W) tag = f'{split}/{k}' if self.for_specs: # flipping values along frequency dim, otherwise mels are upside-down (1, F, T) grid = torchvision.utils.make_grid(images[k].flip(dims=(2,)), nrow=1) # also reconstruct waveform given the spec and inv_transform if k not in ['conditioning', 'conditioning_rec', 'att_nopix', 'att_half', 'att_det']: self._log_rec_audio(images[k], tag, pl_module.global_step, pl_module=pl_module) else: grid = torchvision.utils.make_grid(images[k])# (B,C=1 or 3,H,W) -> (C=3,B*H,W) # attention is already in [0, 1] therefore ignoring this line fig = plt.figure() plt.pcolor(grid.mean(dim=0),vmin=self.melvmin,vmax=self.melvmax) pl_module.logger.experiment.add_figure(tag, fig,global_step=pl_module.global_step) @rank_zero_only def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): root = os.path.join(save_dir, "images", split) for k in images: grid = torchvision.utils.make_grid(images[k], nrow=4) grid = grid.mean(dim=0) grid = grid.numpy() filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( k, global_step, current_epoch, batch_idx) path = os.path.join(root, filename) os.makedirs(os.path.split(path)[0], exist_ok=True) plt.imsave(path,grid,vmin=self.melvmin,vmax=self.melvmax) # also save audio on the disk if self.for_specs: tag = f'{split}/{k}' filename = filename.replace('.png', '.wav') path = os.path.join(root, filename) if k not in ['conditioning', 'conditioning_rec', 'att_nopix', 'att_half', 'att_det']: self._log_rec_audio(images[k], tag, global_step, save_rec_path=path) class CUDACallback(Callback): # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py def on_train_epoch_start(self, trainer, pl_module): # Reset the memory use counter torch.cuda.reset_peak_memory_stats(trainer.strategy.root_device.index) torch.cuda.synchronize(trainer.strategy.root_device.index) self.start_time = time.time() def on_train_epoch_end(self, trainer, pl_module):# ,outputs： outputs positional argument has been removed in the later pytorch-lighning version。 torch.cuda.synchronize(trainer.strategy.root_device.index) max_memory = torch.cuda.max_memory_allocated(trainer.strategy.root_device.index) / 2 ** 20 epoch_time = time.time() - self.start_time try: max_memory = trainer.strategy.reduce(max_memory) epoch_time = trainer.strategy.reduce(epoch_time) rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds") rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB") except AttributeError: pass if __name__ == "__main__": now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") sys.path.append(os.getcwd()) parser = get_parser() parser = Trainer.add_argparse_args(parser) opt, unknown = parser.parse_known_args() if opt.name and opt.resume: raise ValueError( "-n/--name and -r/--resume cannot be specified both." "If you want to resume training in a new log folder, " "use -n/--name in combination with --resume_from_checkpoint" ) if opt.resume: if not os.path.exists(opt.resume): raise ValueError("Cannot find {}".format(opt.resume)) if os.path.isfile(opt.resume): paths = opt.resume.split("/") logdir = "/".join(paths[:-2]) ckpt = opt.resume else: assert os.path.isdir(opt.resume), opt.resume logdir = opt.resume.rstrip("/") ckpt = os.path.join(logdir, "checkpoints", "last.ckpt") opt.ckpt_path = ckpt base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml"))) opt.base = base_configs + opt.base _tmp = logdir.split("/") nowname = _tmp[-1] else: if opt.name: name = "_" + opt.name elif opt.base: cfg_fname = os.path.split(opt.base[0])[-1] cfg_name = os.path.splitext(cfg_fname)[0] name = "_" + cfg_name else: name = "" nowname = now + name + opt.postfix logdir = os.path.join(opt.logdir, nowname) ckptdir = os.path.join(logdir, "checkpoints") cfgdir = os.path.join(logdir, "configs") seed_everything(opt.seed) try: # init and save configs print(f"opt.base:{opt.base}") configs = [OmegaConf.load(cfg) for cfg in opt.base] cli = OmegaConf.from_dotlist(unknown) config = OmegaConf.merge(*configs, cli) lightning_config = config.pop("lightning", OmegaConf.create()) # merge trainer cli with config trainer_config = lightning_config.get("trainer", OmegaConf.create()) # default to ddp trainer_config["strategy"] = "ddp" # "ddp" # "ddp_find_unused_parameters_false" for k in nondefault_trainer_args(opt): trainer_config[k] = getattr(opt, k) if not "gpus" in trainer_config: del trainer_config["strategy"] cpu = True else: gpuinfo = trainer_config["gpus"] print(f"Running on GPUs {gpuinfo}") cpu = False trainer_opt = argparse.Namespace(**trainer_config) lightning_config.trainer = trainer_config # model model = instantiate_from_config(config.model) # trainer and callbacks trainer_kwargs = dict() # default logger configs default_logger_cfgs = { "wandb": { "target": "pytorch_lightning.loggers.WandbLogger", "params": { "name": nowname, "save_dir": logdir, "offline": opt.debug, "id": nowname, } }, "tensorboard": { "target": "pytorch_lightning.loggers.TensorBoardLogger", "params": { "name": "tensorboard", "save_dir": logdir, } }, } default_logger_cfg = default_logger_cfgs["tensorboard"] if "logger" in lightning_config: logger_cfg = lightning_config.logger else: logger_cfg = OmegaConf.create() logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg) trainer_kwargs["logger"] = instantiate_from_config(logger_cfg) default_modelckpt_cfg = { "target": "pytorch_lightning.callbacks.ModelCheckpoint", "params": { "dirpath": ckptdir, "filename": "{epoch:06}", "verbose": True, "save_last": True, "save_top_k": 5, } } # use valitdation monitor: if hasattr(model, "monitor"): print(f"Monitoring {model.monitor} as checkpoint metric.") default_modelckpt_cfg["params"]["monitor"] = model.monitor if "modelcheckpoint" in lightning_config: modelckpt_cfg = lightning_config.modelcheckpoint else: modelckpt_cfg = OmegaConf.create() modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg) print(f"Merged modelckpt-cfg: \n{modelckpt_cfg}") # add callback which sets up log directory default_callbacks_cfg = { "setup_callback": { "target": "main.SetupCallback", "params": { "resume": opt.resume, "now": now, "logdir": logdir, "ckptdir": ckptdir, "cfgdir": cfgdir, "config": config, "lightning_config": lightning_config, } }, "image_logger": { "target": "main.ImageLogger", "params": { "batch_frequency": 5000, "max_images": 4, } }, "learning_rate_logger": { "target": "main.LearningRateMonitor", "params": { "logging_interval": "step", # "log_momentum": True } }, "cuda_callback": { "target": "main.CUDACallback" }, } # patching the default config for the spectrogram input # if 'Spectrogram' in config.data.target: # spec_dir_name = Path(config.data.params.spec_dir_path).name # default_callbacks_cfg['image_logger']['params']['spec_dir_name'] = spec_dir_name # default_callbacks_cfg['image_logger']['params']['sample_rate'] = config.data.params.sample_rate default_callbacks_cfg.update({'checkpoint_callback': modelckpt_cfg}) if "callbacks" in lightning_config: callbacks_cfg = lightning_config.callbacks else: callbacks_cfg = OmegaConf.create() if 'metrics_over_trainsteps_checkpoint' in callbacks_cfg: print( 'Caution: Saving checkpoints every n train steps without deleting. This might require some free space.') default_metrics_over_trainsteps_ckpt_dict = { 'metrics_over_trainsteps_checkpoint': {"target": 'pytorch_lightning.callbacks.ModelCheckpoint', 'params': { "dirpath": os.path.join(ckptdir, 'trainstep_checkpoints'), "filename": "{epoch:06}-{step:09}", "verbose": True, 'save_top_k': -1, 'every_n_train_steps': 10000, 'save_weights_only': True } } } default_callbacks_cfg.update(default_metrics_over_trainsteps_ckpt_dict) callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg) if 'ignore_keys_callback' in callbacks_cfg and hasattr(trainer_opt, 'ckpt_path'):# false for the former callbacks_cfg.ignore_keys_callback.params['ckpt_path'] = trainer_opt.ckpt_path elif 'ignore_keys_callback' in callbacks_cfg: del callbacks_cfg['ignore_keys_callback'] trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg] trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs) trainer.logdir = logdir ##### data ##### data = instantiate_from_config(config.data) data.prepare_data() data.setup() print("#### Data #####") for k in data.datasets: print(f"{k}, {data.datasets[k].__class__.__name__}, {len(data.datasets[k])}") # configure learning rate bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate if not cpu: ngpu = len(lightning_config.trainer.gpus.strip(",").split(',')) else: ngpu = 1 if 'accumulate_grad_batches' in lightning_config.trainer: accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches else: accumulate_grad_batches = 1 print(f"accumulate_grad_batches = {accumulate_grad_batches}") lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches if opt.scale_lr: model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr print( "Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (base_lr)".format( model.learning_rate, accumulate_grad_batches, ngpu, bs, base_lr)) else: model.learning_rate = base_lr print("++++ NOT USING LR SCALING ++++") print(f"Setting learning rate to {model.learning_rate:.2e}") # allow checkpointing via USR1 def melk(*args, **kwargs): # run all checkpoint hooks if trainer.global_rank == 0: print("Summoning checkpoint.") ckpt_path = os.path.join(ckptdir, "last.ckpt") trainer.save_checkpoint(ckpt_path) def divein(*args, **kwargs): if trainer.global_rank == 0: import pudb; pudb.set_trace() import signal signal.signal(signal.SIGUSR1, melk) signal.signal(signal.SIGUSR2, divein) print(f"##### trainer.logdir:{trainer.logdir} #####") # run if opt.train: try: if hasattr(opt,'ckpt_path'): trainer.fit(model, data,ckpt_path = opt.ckpt_path) else: trainer.fit(model, data) except Exception: melk() raise elif opt.val: trainer.validate(model, data) if not opt.no_test and not trainer.interrupted: if not opt.train and hasattr(opt,'ckpt_path'):# just test the ckeckpoint, without training trainer.test(model, data, ckpt_path = opt.ckpt_path) else:# test the model after trainning trainer.test(model, data) except Exception: if opt.debug and trainer.global_rank == 0: try: import pudb as debugger except ImportError: import pdb as debugger debugger.post_mortem() raise finally: # move newly created debug project to debug_runs if opt.debug and not opt.resume and trainer.global_rank == 0: dst, name = os.path.split(logdir) dst = os.path.join(dst, "debug_runs", name) os.makedirs(os.path.split(dst)[0], exist_ok=True) os.rename(logdir, dst) if trainer.global_rank == 0: print(trainer.profiler.summary())