Spaces:
Configuration error
Configuration error
import argparse, os, sys, datetime, glob | |
import numpy as np | |
import time | |
import torch | |
import torchvision | |
import pytorch_lightning as pl | |
import json | |
import pickle | |
from packaging import version | |
from omegaconf import OmegaConf | |
from torch.utils.data import DataLoader, Dataset | |
from functools import partial | |
from PIL import Image | |
import torch.distributed as dist | |
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 pytorch_lightning.plugins import DDPPlugin | |
sys.path.append("./stable_diffusion") | |
from ldm.data.base import Txt2ImgIterableBaseDataset | |
from ldm.util import instantiate_from_config | |
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( | |
"--no-test", | |
type=str2bool, | |
const=True, | |
default=False, | |
nargs="?", | |
help="disable 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", | |
action="store_true", | |
default=False, | |
help="scale base-lr by ngpu * batch_size * n_accumulate", | |
) | |
return parser | |
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 | |
if isinstance(dataset, Txt2ImgIterableBaseDataset): | |
split_size = dataset.num_records // worker_info.num_workers | |
# reset num_records to the true number to retain reliable length information | |
dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size] | |
current_id = np.random.choice(len(np.random.get_state()[1]), 1) | |
return np.random.seed(np.random.get_state()[1][current_id] + worker_id) | |
else: | |
return np.random.seed(np.random.get_state()[1][0] + worker_id) | |
class DataModuleFromConfig(pl.LightningDataModule): | |
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): | |
is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset) | |
if is_iterable_dataset or self.use_worker_init_fn: | |
init_fn = worker_init_fn | |
else: | |
init_fn = None | |
return DataLoader(self.datasets["train"], batch_size=self.batch_size, | |
num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True, | |
worker_init_fn=init_fn, persistent_workers=True) | |
def _val_dataloader(self, shuffle=False): | |
if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn: | |
init_fn = worker_init_fn | |
else: | |
init_fn = None | |
return DataLoader(self.datasets["validation"], | |
batch_size=self.batch_size, | |
num_workers=self.num_workers, | |
worker_init_fn=init_fn, | |
shuffle=shuffle, persistent_workers=True) | |
def _test_dataloader(self, shuffle=False): | |
is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset) | |
if is_iterable_dataset or self.use_worker_init_fn: | |
init_fn = worker_init_fn | |
else: | |
init_fn = None | |
# do not shuffle dataloader for iterable dataset | |
shuffle = shuffle and (not is_iterable_dataset) | |
return DataLoader(self.datasets["test"], batch_size=self.batch_size, | |
num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle, persistent_workers=True) | |
def _predict_dataloader(self, shuffle=False): | |
if isinstance(self.datasets['predict'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn: | |
init_fn = worker_init_fn | |
else: | |
init_fn = None | |
return DataLoader(self.datasets["predict"], batch_size=self.batch_size, | |
num_workers=self.num_workers, worker_init_fn=init_fn, persistent_workers=True) | |
class SetupCallback(Callback): | |
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_keyboard_interrupt(self, trainer, pl_module): | |
if trainer.global_rank == 0: | |
print("Summoning checkpoint.") | |
ckpt_path = os.path.join(self.ckptdir, "last.ckpt") | |
trainer.save_checkpoint(ckpt_path) | |
def on_pretrain_routine_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))) | |
def get_world_size(): | |
if not dist.is_available(): | |
return 1 | |
if not dist.is_initialized(): | |
return 1 | |
return dist.get_world_size() | |
def all_gather(data): | |
""" | |
Run all_gather on arbitrary picklable data (not necessarily tensors) | |
Args: | |
data: any picklable object | |
Returns: | |
list[data]: list of data gathered from each rank | |
""" | |
world_size = get_world_size() | |
if world_size == 1: | |
return [data] | |
# serialized to a Tensor | |
origin_size = None | |
if not isinstance(data, torch.Tensor): | |
buffer = pickle.dumps(data) | |
storage = torch.ByteStorage.from_buffer(buffer) | |
tensor = torch.ByteTensor(storage).to("cuda") | |
else: | |
origin_size = data.size() | |
tensor = data.reshape(-1) | |
tensor_type = tensor.dtype | |
# obtain Tensor size of each rank | |
local_size = torch.LongTensor([tensor.numel()]).to("cuda") | |
size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)] | |
dist.all_gather(size_list, local_size) | |
size_list = [int(size.item()) for size in size_list] | |
max_size = max(size_list) | |
# receiving Tensor from all ranks | |
# we pad the tensor because torch all_gather does not support | |
# gathering tensors of different shapes | |
tensor_list = [] | |
for _ in size_list: | |
tensor_list.append(torch.FloatTensor(size=(max_size,)).cuda().to(tensor_type)) | |
if local_size != max_size: | |
padding = torch.FloatTensor(size=(max_size - local_size,)).cuda().to(tensor_type) | |
tensor = torch.cat((tensor, padding), dim=0) | |
dist.all_gather(tensor_list, tensor) | |
data_list = [] | |
for size, tensor in zip(size_list, tensor_list): | |
if origin_size is None: | |
buffer = tensor.cpu().numpy().tobytes()[:size] | |
data_list.append(pickle.loads(buffer)) | |
else: | |
buffer = tensor[:size] | |
data_list.append(buffer) | |
if origin_size is not None: | |
new_shape = [-1] + list(origin_size[1:]) | |
resized_list = [] | |
for data in data_list: | |
# suppose the difference of tensor size exist in first dimension | |
data = data.reshape(new_shape) | |
resized_list.append(data) | |
return resized_list | |
else: | |
return data_list | |
class ImageLogger(Callback): | |
def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True, | |
rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False, | |
log_images_kwargs=None): | |
super().__init__() | |
self.rescale = rescale | |
self.batch_freq = batch_frequency | |
self.max_images = max_images | |
self.logger_log_images = { | |
pl.loggers.TestTubeLogger: self._testtube, | |
} | |
self.log_steps = [2 ** n for n in range(6, int(np.log2(self.batch_freq)) + 1)] | |
if not increase_log_steps: | |
self.log_steps = [self.batch_freq] | |
self.clamp = clamp | |
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 | |
def _testtube(self, pl_module, images, batch_idx, split): | |
for k in images: | |
grid = torchvision.utils.make_grid(images[k]) | |
grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w | |
tag = f"{split}/{k}" | |
pl_module.logger.experiment.add_image( | |
tag, grid, | |
global_step=pl_module.global_step) | |
def log_local(self, save_dir, split, images, prompts, | |
global_step, current_epoch, batch_idx): | |
root = os.path.join(save_dir, "images", split) | |
names = {"reals": "before", "inputs": "after", "reconstruction": "before-vq", "samples": "after-gen"} | |
# print(root) | |
for k in images: | |
grid = torchvision.utils.make_grid(images[k], nrow=8) | |
if self.rescale: | |
grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w | |
grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) | |
grid = grid.numpy() | |
grid = (grid * 255).astype(np.uint8) | |
filename = "gs-{:06}_e-{:06}_b-{:06}_{}.png".format( | |
global_step, | |
current_epoch, | |
batch_idx, | |
names[k]) | |
path = os.path.join(root, filename) | |
os.makedirs(os.path.split(path)[0], exist_ok=True) | |
# print(path) | |
Image.fromarray(grid).save(path) | |
filename = "gs-{:06}_e-{:06}_b-{:06}_prompt.json".format( | |
global_step, | |
current_epoch, | |
batch_idx) | |
path = os.path.join(root, filename) | |
with open(path, "w") as f: | |
for p in prompts: | |
f.write(f"{json.dumps(p)}\n") | |
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) or (split == "val" and batch_idx == 0): | |
logger = type(pl_module.logger) | |
is_train = pl_module.training | |
if is_train: | |
pl_module.eval() | |
with torch.no_grad(): | |
images = pl_module.log_images(batch, split=split, **self.log_images_kwargs) | |
prompts = batch["edit"]["c_crossattn"][:self.max_images] | |
prompts = [p for ps in all_gather(prompts) for p in ps] | |
for k in images: | |
N = min(images[k].shape[0], self.max_images) | |
images[k] = images[k][:N] | |
images[k] = torch.cat(all_gather(images[k][:N])) | |
if isinstance(images[k], torch.Tensor): | |
images[k] = images[k].detach().cpu() | |
if self.clamp: | |
images[k] = torch.clamp(images[k], -1., 1.) | |
self.log_local(pl_module.logger.save_dir, split, images, prompts, | |
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): | |
if len(self.log_steps) > 0: | |
self.log_steps.pop(0) | |
return True | |
return False | |
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_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") | |
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 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.root_gpu) | |
torch.cuda.synchronize(trainer.root_gpu) | |
self.start_time = time.time() | |
def on_train_epoch_end(self, trainer, pl_module, outputs): | |
torch.cuda.synchronize(trainer.root_gpu) | |
max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20 | |
epoch_time = time.time() - self.start_time | |
try: | |
max_memory = trainer.training_type_plugin.reduce(max_memory) | |
epoch_time = trainer.training_type_plugin.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__": | |
# custom parser to specify config files, train, test and debug mode, | |
# postfix, resume. | |
# `--key value` arguments are interpreted as arguments to the trainer. | |
# `nested.key=value` arguments are interpreted as config parameters. | |
# configs are merged from left-to-right followed by command line parameters. | |
# model: | |
# base_learning_rate: float | |
# target: path to lightning module | |
# params: | |
# key: value | |
# data: | |
# target: main.DataModuleFromConfig | |
# params: | |
# batch_size: int | |
# wrap: bool | |
# train: | |
# target: path to train dataset | |
# params: | |
# key: value | |
# validation: | |
# target: path to validation dataset | |
# params: | |
# key: value | |
# test: | |
# target: path to test dataset | |
# params: | |
# key: value | |
# lightning: (optional, has sane defaults and can be specified on cmdline) | |
# trainer: | |
# additional arguments to trainer | |
# logger: | |
# logger to instantiate | |
# modelcheckpoint: | |
# modelcheckpoint to instantiate | |
# callbacks: | |
# callback1: | |
# target: importpath | |
# params: | |
# key: value | |
now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") | |
# add cwd for convenience and to make classes in this file available when | |
# running as `python main.py` | |
# (in particular `main.DataModuleFromConfig`) | |
sys.path.append(os.getcwd()) | |
parser = get_parser() | |
parser = Trainer.add_argparse_args(parser) | |
opt, unknown = parser.parse_known_args() | |
assert opt.name | |
cfg_fname = os.path.split(opt.base[0])[-1] | |
cfg_name = os.path.splitext(cfg_fname)[0] | |
nowname = f"{cfg_name}_{opt.name}" | |
logdir = os.path.join(opt.logdir, nowname) | |
ckpt = os.path.join(logdir, "checkpoints", "last.ckpt") | |
resume = False | |
if os.path.isfile(ckpt): | |
opt.resume_from_checkpoint = ckpt | |
base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml"))) | |
opt.base = base_configs + opt.base | |
_tmp = logdir.split("/") | |
nowname = _tmp[-1] | |
resume = True | |
ckptdir = os.path.join(logdir, "checkpoints") | |
cfgdir = os.path.join(logdir, "configs") | |
os.makedirs(logdir, exist_ok=True) | |
os.makedirs(ckptdir, exist_ok=True) | |
os.makedirs(cfgdir, exist_ok=True) | |
try: | |
# init and save configs | |
configs = [OmegaConf.load(cfg) for cfg in opt.base] | |
cli = OmegaConf.from_dotlist(unknown) | |
config = OmegaConf.merge(*configs, cli) | |
if resume: | |
# By default, when finetuning from Stable Diffusion, we load the EMA-only checkpoint to initialize all weights. | |
# If resuming InstructPix2Pix from a finetuning checkpoint, instead load both EMA and non-EMA weights. | |
config.model.params.load_ema = True | |
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["accelerator"] = "ddp" | |
for k in nondefault_trainer_args(opt): | |
trainer_config[k] = getattr(opt, k) | |
if not "gpus" in trainer_config: | |
del trainer_config["accelerator"] | |
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, | |
"id": nowname, | |
} | |
}, | |
"testtube": { | |
"target": "pytorch_lightning.loggers.TestTubeLogger", | |
"params": { | |
"name": "testtube", | |
"save_dir": logdir, | |
} | |
}, | |
} | |
default_logger_cfg = default_logger_cfgs["wandb"] | |
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) | |
# modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to | |
# specify which metric is used to determine best models | |
default_modelckpt_cfg = { | |
"target": "pytorch_lightning.callbacks.ModelCheckpoint", | |
"params": { | |
"dirpath": ckptdir, | |
"filename": "{epoch:06}", | |
"verbose": True, | |
"save_last": True, | |
} | |
} | |
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}") | |
if version.parse(pl.__version__) < version.parse('1.4.0'): | |
trainer_kwargs["checkpoint_callback"] = instantiate_from_config(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": 750, | |
"max_images": 4, | |
"clamp": True | |
} | |
}, | |
"learning_rate_logger": { | |
"target": "main.LearningRateMonitor", | |
"params": { | |
"logging_interval": "step", | |
# "log_momentum": True | |
} | |
}, | |
"cuda_callback": { | |
"target": "main.CUDACallback" | |
}, | |
} | |
if version.parse(pl.__version__) >= version.parse('1.4.0'): | |
default_callbacks_cfg.update({'checkpoint_callback': modelckpt_cfg}) | |
if "callbacks" in lightning_config: | |
callbacks_cfg = lightning_config.callbacks | |
else: | |
callbacks_cfg = OmegaConf.create() | |
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': 1000, | |
'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, 'resume_from_checkpoint'): | |
callbacks_cfg.ignore_keys_callback.params['ckpt_path'] = trainer_opt.resume_from_checkpoint | |
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, plugins=DDPPlugin(find_unused_parameters=False), **trainer_kwargs) | |
trainer.logdir = logdir ### | |
# data | |
data = instantiate_from_config(config.data) | |
# NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html | |
# calling these ourselves should not be necessary but it is. | |
# lightning still takes care of proper multiprocessing though | |
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) | |
# run | |
if opt.train: | |
try: | |
trainer.fit(model, data) | |
except Exception: | |
melk() | |
raise | |
if not opt.no_test and not trainer.interrupted: | |
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()) | |