Spaces:
Running
Running
| """This code is taken from <https://github.com/alexandre01/deepsvg> | |
| by Alexandre Carlier, Martin Danelljan, Alexandre Alahi and Radu Timofte | |
| from the paper >https://arxiv.org/pdf/2007.11301.pdf> | |
| """ | |
| import torch.optim as optim | |
| from src.preprocessing.deepsvg.deepsvg_schedulers.warmup import GradualWarmupScheduler | |
| class _Config: | |
| """ | |
| Training config. | |
| """ | |
| def __init__(self, num_gpus=1): | |
| self.num_gpus = num_gpus # | |
| self.dataloader_module = "deepsvg.svgtensor_dataset" # | |
| self.collate_fn = None # | |
| self.data_dir = "./data/svgs_tensors/" # | |
| self.meta_filepath = "./data/svgs_meta.csv" # | |
| self.loader_num_workers = 0 # | |
| self.pretrained_path = "./models/hierarchical_ordered.pth.tar" # | |
| self.model_cfg = None # | |
| self.num_epochs = None # | |
| self.num_steps = None # | |
| self.learning_rate = 1e-3 # | |
| self.batch_size = 100 # | |
| self.warmup_steps = 500 # | |
| # Dataset | |
| self.train_ratio = 1.0 # | |
| self.nb_augmentations = 1 # | |
| self.max_num_groups = 15 # | |
| self.max_seq_len = 30 # | |
| self.max_total_len = None # | |
| self.filter_uni = None # | |
| self.filter_category = None # | |
| self.filter_platform = None # | |
| self.filter_labels = None # | |
| self.grad_clip = None # | |
| self.log_every = 20 # | |
| self.val_every = 1000 # | |
| self.ckpt_every = 1000 # | |
| self.stats_to_print = { | |
| "train": ["lr", "time"] | |
| } | |
| self.model_args = [] # | |
| self.optimizer_starts = [0] # | |
| # Overridable methods | |
| def make_model(self): | |
| raise NotImplementedError | |
| def make_losses(self): | |
| raise NotImplementedError | |
| def make_optimizers(self, model): | |
| return [optim.AdamW(model.parameters(), self.learning_rate)] | |
| def make_schedulers(self, optimizers, epoch_size): | |
| return [None] * len(optimizers) | |
| def make_warmup_schedulers(self, optimizers, scheduler_lrs): | |
| return [GradualWarmupScheduler(optimizer, multiplier=1.0, total_epoch=self.warmup_steps, after_scheduler=scheduler_lr) | |
| for optimizer, scheduler_lr in zip(optimizers, scheduler_lrs)] | |
| def get_params(self, step, epoch): | |
| return {} | |
| def get_weights(self, step, epoch): | |
| return {} | |
| def set_train_vars(self, train_vars, dataloader): | |
| pass | |
| def visualize(self, model, output, train_vars, step, epoch, summary_writer, visualization_dir): | |
| pass | |
| # Utility methods | |
| def values(self): | |
| for key in dir(self): | |
| if not key.startswith("__") and not callable(getattr(self, key)): | |
| yield key, getattr(self, key) | |
| def to_dict(self): | |
| return {key: val for key, val in self.values()} | |
| def load_dict(self, dict): | |
| for key, val in dict.items(): | |
| setattr(self, key, val) | |
| def print_params(self): | |
| for key, val in self.values(): | |
| print(f" {key} = {val}") | |