# -*- coding: utf-8 -*- # Copyright (c) XiMing Xing. All rights reserved. # Author: XiMing Xing # Description: import math import copy import random import pathlib from typing import Dict from shapely.geometry.polygon import Polygon import omegaconf import cv2 import numpy as np import pydiffvg import torch from torch.optim.lr_scheduler import LambdaLR from pytorch_svgrender.diffvg_warp import DiffVGState from pytorch_svgrender.libs.solver.optim import get_optimizer class Painter(DiffVGState): def __init__( self, diffvg_cfg: omegaconf.DictConfig, style: str, num_segments: int, segment_init: str, radius: int = 20, canvas_size: int = 600, n_grid: int = 32, trainable_bg: bool = False, stroke_width: int = 3, path_svg=None, device=None, ): super().__init__(device, print_timing=diffvg_cfg.print_timing, canvas_width=canvas_size, canvas_height=canvas_size) self.style = style self.num_segments = num_segments self.segment_init = segment_init self.radius = radius """pixelart params""" self.n_grid = n_grid # divide the canvas into n grids self.pixel_per_grid = self.canvas_width // self.n_grid """sketch params""" self.stroke_width = stroke_width """iconography params""" self.color_ref = None self.path_svg = path_svg self.optimize_flag = [] self.strokes_counter = 0 # counts the number of calls to "get_path" # Background color self.para_bg = torch.tensor([1., 1., 1.], requires_grad=trainable_bg, device=self.device) self.target_img = None self.pos_init_method = None def component_wise_path_init(self, gt, pred, init_type: str = 'sparse'): # set target image self.target_img = gt if init_type == 'random': self.pos_init_method = RandomCoordInit(self.canvas_height, self.canvas_width) elif init_type == 'sparse': # when initialized for the first time, the render result is None if pred is None: pred = self.para_bg.view(1, -1, 1, 1).repeat(1, 1, self.canvas_height, self.canvas_width) # then pred is the render result self.pos_init_method = SparseCoordInit(pred, gt) elif init_type == 'naive': if pred is None: pred = self.para_bg.view(1, -1, 1, 1).repeat(1, 1, self.canvas_height, self.canvas_width) self.pos_init_method = NaiveCoordInit(pred, gt) else: raise NotImplementedError(f"'{init_type}' is not support.") def init_image(self, stage=0, num_paths=0): self.cur_shapes, self.cur_shape_groups = [], [] # or init svg by pydiffvg if self.style in ['pixelart', 'low-poly']: # update path definition num_paths = self.n_grid if stage > 0: # Noting: if multi stages training than add new strokes on existing ones # don't optimize on previous strokes self.optimize_flag = [False for i in range(len(self.shapes))] for i in range(num_paths): if self.style == 'iconography': path = self.get_path() self.shapes.append(path) self.cur_shapes.append(path) fill_color_init = torch.FloatTensor(np.random.uniform(size=[4])) fill_color_init[-1] = 1.0 path_group = pydiffvg.ShapeGroup( shape_ids=torch.tensor([self.strokes_counter - 1]), fill_color=fill_color_init, stroke_color=None ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) self.optimize_flag.append(True) elif self.style in ['pixelart', 'low-poly']: for j in range(num_paths): path = self.get_path(coord=[i, j]) self.shapes.append(path) self.cur_shapes.append(path) fill_color_init = torch.FloatTensor(np.random.uniform(size=[4])) fill_color_init[-1] = 1.0 path_group = pydiffvg.ShapeGroup( shape_ids=torch.LongTensor([i * num_paths + j]), fill_color=fill_color_init, stroke_color=None, ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) self.optimize_flag.append(True) elif self.style in ['ink', 'sketch']: path = self.get_path() self.shapes.append(path) self.cur_shapes.append(path) stroke_color_init = [0.0, 0.0, 0.0] + [random.random()] stroke_color_init = torch.FloatTensor(stroke_color_init) path_group = pydiffvg.ShapeGroup( shape_ids=torch.tensor([len(self.shapes) - 1]), fill_color=None, stroke_color=stroke_color_init ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) elif self.style == 'painting': path = self.get_path() self.shapes.append(path) self.cur_shapes.append(path) wref, href = self.color_ref wref = max(0, min(int(wref), self.canvas_width - 1)) href = max(0, min(int(href), self.canvas_height - 1)) stroke_color_init = list(self.target_img[0, :, href, wref]) + [1.] path_group = pydiffvg.ShapeGroup( shape_ids=torch.tensor([len(self.shapes) - 1]), fill_color=None, stroke_color=torch.FloatTensor(stroke_color_init) ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) else: num_paths_exists = 0 if self.path_svg is not None and pathlib.Path(self.path_svg).exists(): print(f"-> init svg from `{self.path_svg}` ...") self.canvas_width, self.canvas_height, self.shapes, self.shape_groups = self.load_svg(self.path_svg) # if you want to add more strokes to existing ones and optimize on all of them num_paths_exists = len(self.shapes) self.cur_shapes = self.shapes self.cur_shape_groups = self.shape_groups for i in range(num_paths_exists, num_paths): if self.style == 'iconography': path = self.get_path() self.shapes.append(path) self.cur_shapes.append(path) wref, href = self.color_ref wref = max(0, min(int(wref), self.canvas_width - 1)) href = max(0, min(int(href), self.canvas_height - 1)) fill_color_init = list(self.target_img[0, :, href, wref]) + [1.] path_group = pydiffvg.ShapeGroup( shape_ids=torch.tensor([self.strokes_counter - 1]), fill_color=torch.FloatTensor(fill_color_init), stroke_color=None ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) elif self.style in ['pixelart', 'low-poly']: for j in range(num_paths): path = self.get_path(coord=[i, j]) self.shapes.append(path) self.cur_shapes.append(path) fill_color_init = torch.FloatTensor(np.random.uniform(size=[4])) fill_color_init[-1] = 1.0 path_group = pydiffvg.ShapeGroup( shape_ids=torch.LongTensor([i * num_paths + j]), fill_color=fill_color_init, stroke_color=None, ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) elif self.style in ['sketch', 'ink']: path = self.get_path() self.shapes.append(path) self.cur_shapes.append(path) stroke_color_init = [0.0, 0.0, 0.0] + [random.random()] stroke_color_init = torch.FloatTensor(stroke_color_init) path_group = pydiffvg.ShapeGroup( shape_ids=torch.tensor([len(self.shapes) - 1]), fill_color=None, stroke_color=stroke_color_init ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) elif self.style in ['painting']: path = self.get_path() self.shapes.append(path) self.cur_shapes.append(path) if self.color_ref is None: stroke_color_val = np.random.uniform(size=[4]) stroke_color_val[-1] = 1.0 stroke_color_init = torch.FloatTensor(stroke_color_val) else: wref, href = self.color_ref wref = max(0, min(int(wref), self.canvas_width - 1)) href = max(0, min(int(href), self.canvas_height - 1)) stroke_color_init = list(self.target_img[0, :, href, wref]) + [1.] stroke_color_init = torch.FloatTensor(stroke_color_init) path_group = pydiffvg.ShapeGroup( shape_ids=torch.tensor([len(self.shapes) - 1]), fill_color=None, stroke_color=stroke_color_init ) self.shape_groups.append(path_group) self.cur_shape_groups.append(path_group) self.optimize_flag = [True for i in range(len(self.shapes))] img = self.get_image() return img def get_image(self, step: int = 0): img = self.render_warp(step) img = img[:, :, 3:4] * img[:, :, :3] + self.para_bg * (1 - img[:, :, 3:4]) img = img.unsqueeze(0) # convert img from HWC to NCHW img = img.permute(0, 3, 1, 2).to(self.device) # NHWC -> NCHW return img def get_path(self, coord=None): num_segments = self.num_segments points = [] if self.style == 'iconography': # init segment if self.segment_init == 'circle': num_control_points = [2] * num_segments radius = self.radius if self.radius is not None else np.random.uniform(0.5, 1) if self.pos_init_method is not None: center = self.pos_init_method() else: center = (random.random(), random.random()) bias = center self.color_ref = copy.deepcopy(bias) avg_degree = 360 / (num_segments * 3) for i in range(0, num_segments * 3): point = ( np.cos(np.deg2rad(i * avg_degree)), np.sin(np.deg2rad(i * avg_degree)) ) points.append(point) points = torch.FloatTensor(points) * radius + torch.FloatTensor(bias).unsqueeze(dim=0) elif self.segment_init == 'random': num_control_points = [2] * num_segments p0 = self.pos_init_method() self.color_ref = copy.deepcopy(p0) points.append(p0) for j in range(num_segments): radius = self.radius p1 = (p0[0] + radius * np.random.uniform(-0.5, 0.5), p0[1] + radius * np.random.uniform(-0.5, 0.5)) p2 = (p1[0] + radius * np.random.uniform(-0.5, 0.5), p1[1] + radius * np.random.uniform(-0.5, 0.5)) p3 = (p2[0] + radius * np.random.uniform(-0.5, 0.5), p2[1] + radius * np.random.uniform(-0.5, 0.5)) points.append(p1) points.append(p2) if j < num_segments - 1: points.append(p3) p0 = p3 points = torch.FloatTensor(points) else: raise NotImplementedError(f"{self.segment_init} is not exists.") path = pydiffvg.Path( num_control_points=torch.LongTensor(num_control_points), points=points, stroke_width=torch.tensor(0.0), is_closed=True ) elif self.style in ['sketch', 'painting', 'ink']: num_control_points = torch.zeros(num_segments, dtype=torch.long) + 2 points = [] p0 = [random.random(), random.random()] points.append(p0) # select color by first point coordinate color_ref = copy.deepcopy(p0) color_ref[0] *= self.canvas_width color_ref[1] *= self.canvas_height self.color_ref = color_ref for j in range(num_segments): radius = 0.1 p1 = (p0[0] + radius * (random.random() - 0.5), p0[1] + radius * (random.random() - 0.5)) p2 = (p1[0] + radius * (random.random() - 0.5), p1[1] + radius * (random.random() - 0.5)) p3 = (p2[0] + radius * (random.random() - 0.5), p2[1] + radius * (random.random() - 0.5)) points.append(p1) points.append(p2) points.append(p3) p0 = p3 points = torch.tensor(points).to(self.device) points[:, 0] *= self.canvas_width points[:, 1] *= self.canvas_height path = pydiffvg.Path(num_control_points=torch.LongTensor(num_control_points), points=points, stroke_width=torch.tensor(float(self.stroke_width)), is_closed=False) elif self.style in ['pixelart', 'low-poly']: x = coord[0] * self.pixel_per_grid y = coord[1] * self.pixel_per_grid points = torch.FloatTensor([ [x, y], [x + self.pixel_per_grid, y], [x + self.pixel_per_grid, y + self.pixel_per_grid], [x, y + self.pixel_per_grid] ]).to(self.device) path = pydiffvg.Polygon(points=points, stroke_width=torch.tensor(0.0), is_closed=True) self.strokes_counter += 1 return path def clip_curve_shape(self): if self.style in ['sketch', 'ink']: for group in self.shape_groups: group.stroke_color.data[:3].clamp_(0., 0.) # to force black stroke group.stroke_color.data[-1].clamp_(0., 1.) # clip alpha else: for group in self.shape_groups: if group.stroke_color is not None: group.stroke_color.data.clamp_(0.0, 1.0) # clip rgba if group.fill_color is not None: group.fill_color.data.clamp_(0.0, 1.0) # clip rgba def reinitialize_paths(self, reinit_path: bool = False, opacity_threshold: float = None, area_threshold: float = None, fpath: pathlib.Path = None): """ reinitialize paths, also known as 'Reinitializing paths' in VectorFusion paper. Args: reinit_path: whether to reinitialize paths or not. opacity_threshold: Threshold of opacity. area_threshold: Threshold of the closed polygon area. fpath: The path to save the reinitialized SVG. """ if not reinit_path: return if self.style not in ['iconography', 'low-poly', 'painting']: return def get_keys_below_threshold(my_dict, threshold): keys_below_threshold = [key for key, value in my_dict.items() if value < threshold] return keys_below_threshold select_path_ids_by_opc = [] select_path_ids_by_area = [] if self.style in ['iconography', 'low-poly']: # re-init by opacity_threshold if opacity_threshold != 0 and opacity_threshold is not None: opacity_record_ = {group.shape_ids.item(): group.fill_color[-1].item() for group in self.cur_shape_groups} # print("-> opacity_record: ", opacity_record_) print("-> opacity_record: ", [f"{k}: {v:.3f}" for k, v in opacity_record_.items()]) select_path_ids_by_opc = get_keys_below_threshold(opacity_record_, opacity_threshold) print("select_path_ids_by_opc: ", select_path_ids_by_opc) # remove path by area_threshold if area_threshold != 0 and area_threshold is not None: area_records = [Polygon(shape.points.detach().cpu().numpy()).area for shape in self.cur_shapes] # print("-> area_records: ", area_records) print("-> area_records: ", ['%.2f' % i for i in area_records]) for i, shape in enumerate(self.cur_shapes): points_ = shape.points.detach().cpu().numpy() if Polygon(points_).area < area_threshold: select_path_ids_by_area.append(shape.id) print("select_path_ids_by_area: ", select_path_ids_by_area) elif self.style in ['painting']: # re-init by opacity_threshold if opacity_threshold != 0 and opacity_threshold is not None: opacity_record_ = {group.shape_ids.item(): group.stroke_color[-1].item() for group in self.cur_shape_groups} # print("-> opacity_record: ", opacity_record_) print("-> opacity_record: ", [f"{k}: {v:.3f}" for k, v in opacity_record_.items()]) select_path_ids_by_opc = get_keys_below_threshold(opacity_record_, opacity_threshold) print("select_path_ids_by_opc: ", select_path_ids_by_opc) # re-init paths reinit_union = list(set(select_path_ids_by_opc + select_path_ids_by_area)) if len(reinit_union) > 0: for i, path in enumerate(self.cur_shapes): if path.id in reinit_union: coord = [i, i] if self.style == 'low-poly' else None self.cur_shapes[i] = self.get_path(coord=coord) for i, group in enumerate(self.cur_shape_groups): shp_ids = group.shape_ids.cpu().numpy().tolist() if set(shp_ids).issubset(reinit_union): if self.style in ['iconography', 'low-poly']: fill_color_init = torch.FloatTensor(np.random.uniform(size=[4])) fill_color_init[-1] = 1.0 self.cur_shape_groups[i] = pydiffvg.ShapeGroup( shape_ids=torch.tensor(list(shp_ids)), fill_color=fill_color_init, stroke_color=None) elif self.style in ['painting']: stroke_color_init = torch.FloatTensor(np.random.uniform(size=[4])) stroke_color_init[-1] = 1.0 self.cur_shape_groups[i] = pydiffvg.ShapeGroup( shape_ids=torch.tensor([len(self.shapes) - 1]), fill_color=None, stroke_color=stroke_color_init ) # save reinit svg self.pretty_save_svg(fpath) print("-" * 40) def calc_distance_weight(self, loss_weight_keep): shapes_forsdf = copy.deepcopy(self.cur_shapes) shape_groups_forsdf = copy.deepcopy(self.cur_shape_groups) for si in shapes_forsdf: si.stroke_width = torch.FloatTensor([0]).to(self.device) for sg_idx, sgi in enumerate(shape_groups_forsdf): sgi.fill_color = torch.FloatTensor([1, 1, 1, 1]).to(self.device) sgi.shape_ids = torch.LongTensor([sg_idx]).to(self.device) sargs_forsdf = pydiffvg.RenderFunction.serialize_scene( self.canvas_width, self.canvas_height, shapes_forsdf, shape_groups_forsdf ) _render = pydiffvg.RenderFunction.apply with torch.no_grad(): im_forsdf = _render(self.canvas_width, # width self.canvas_height, # height 2, # num_samples_x 2, # num_samples_y 0, # seed None, *sargs_forsdf) # use alpha channel is a trick to get 0-1 image im_forsdf = (im_forsdf[:, :, 3]).detach().cpu().numpy() loss_weight = get_sdf(im_forsdf, normalize='to1') loss_weight += loss_weight_keep loss_weight = np.clip(loss_weight, 0, 1) loss_weight = torch.FloatTensor(loss_weight).to(self.device) return loss_weight def set_point_parameters(self, id_delta=0): self.point_vars = [] for i, path in enumerate(self.cur_shapes): path.id = i + id_delta # set point id path.points.requires_grad = True self.point_vars.append(path.points) def get_point_parameters(self): return self.point_vars def set_color_parameters(self): self.color_vars = [] for i, group in enumerate(self.cur_shape_groups): if group.fill_color is not None: group.fill_color.requires_grad = True self.color_vars.append(group.fill_color) if group.stroke_color is not None: group.stroke_color.requires_grad = True self.color_vars.append(group.stroke_color) def get_color_parameters(self): return self.color_vars def set_width_parameters(self): # stroke`s width optimization self.width_vars = [] for i, path in enumerate(self.shapes): path.stroke_width.requires_grad = True self.width_vars.append(path.stroke_width) def get_width_parameters(self): return self.width_vars def pretty_save_svg(self, filename, width=None, height=None, shapes=None, shape_groups=None): width = self.canvas_width if width is None else width height = self.canvas_height if height is None else height shapes = self.shapes if shapes is None else shapes shape_groups = self.shape_groups if shape_groups is None else shape_groups self.save_svg(filename, width, height, shapes, shape_groups, use_gamma=False, background=None) def load_svg(self, path_svg): canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(path_svg) return canvas_width, canvas_height, shapes, shape_groups def get_sdf(phi, **kwargs): import skfmm # local import phi = (phi - 0.5) * 2 if (phi.max() <= 0) or (phi.min() >= 0): return np.zeros(phi.shape).astype(np.float32) sd = skfmm.distance(phi, dx=1) flip_negative = kwargs.get('flip_negative', True) if flip_negative: sd = np.abs(sd) truncate = kwargs.get('truncate', 10) sd = np.clip(sd, -truncate, truncate) # print(f"max sd value is: {sd.max()}") zero2max = kwargs.get('zero2max', True) if zero2max and flip_negative: sd = sd.max() - sd elif zero2max: raise ValueError normalize = kwargs.get('normalize', 'sum') if normalize == 'sum': sd /= sd.sum() elif normalize == 'to1': sd /= sd.max() return sd class SparseCoordInit: def __init__(self, pred, gt, format='[bs x c x 2D]', quantile_interval=200, nodiff_thres=0.1): if torch.is_tensor(pred): pred = pred.detach().cpu().numpy() if torch.is_tensor(gt): gt = gt.detach().cpu().numpy() if format == '[bs x c x 2D]': self.map = ((pred[0] - gt[0]) ** 2).sum(0) self.reference_gt = copy.deepcopy(np.transpose(gt[0], (1, 2, 0))) elif format == ['[2D x c]']: self.map = (np.abs(pred - gt)).sum(-1) self.reference_gt = copy.deepcopy(gt[0]) else: raise ValueError # OptionA: Zero too small errors to avoid the error too small deadloop self.map[self.map < nodiff_thres] = 0 quantile_interval = np.linspace(0., 1., quantile_interval) quantized_interval = np.quantile(self.map, quantile_interval) # remove redundant quantized_interval = np.unique(quantized_interval) quantized_interval = sorted(quantized_interval[1:-1]) self.map = np.digitize(self.map, quantized_interval, right=False) self.map = np.clip(self.map, 0, 255).astype(np.uint8) self.idcnt = {} for idi in sorted(np.unique(self.map)): self.idcnt[idi] = (self.map == idi).sum() # remove smallest one to remove the correct region self.idcnt.pop(min(self.idcnt.keys())) def __call__(self): if len(self.idcnt) == 0: h, w = self.map.shape return [np.random.uniform(0, 1) * w, np.random.uniform(0, 1) * h] target_id = max(self.idcnt, key=self.idcnt.get) _, component, cstats, ccenter = cv2.connectedComponentsWithStats( (self.map == target_id).astype(np.uint8), connectivity=4 ) # remove cid = 0, it is the invalid area csize = [ci[-1] for ci in cstats[1:]] target_cid = csize.index(max(csize)) + 1 center = ccenter[target_cid][::-1] coord = np.stack(np.where(component == target_cid)).T dist = np.linalg.norm(coord - center, axis=1) target_coord_id = np.argmin(dist) coord_h, coord_w = coord[target_coord_id] # replace_sampling self.idcnt[target_id] -= max(csize) if self.idcnt[target_id] == 0: self.idcnt.pop(target_id) self.map[component == target_cid] = 0 return [coord_w, coord_h] class RandomCoordInit: def __init__(self, canvas_width, canvas_height): self.canvas_width, self.canvas_height = canvas_width, canvas_height def __call__(self): w, h = self.canvas_width, self.canvas_height return [np.random.uniform(0, 1) * w, np.random.uniform(0, 1) * h] class NaiveCoordInit: def __init__(self, pred, gt, format='[bs x c x 2D]', replace_sampling=True): if isinstance(pred, torch.Tensor): pred = pred.detach().cpu().numpy() if isinstance(gt, torch.Tensor): gt = gt.detach().cpu().numpy() if format == '[bs x c x 2D]': self.map = ((pred[0] - gt[0]) ** 2).sum(0) elif format == ['[2D x c]']: self.map = ((pred - gt) ** 2).sum(-1) else: raise ValueError self.replace_sampling = replace_sampling def __call__(self): coord = np.where(self.map == self.map.max()) coord_h, coord_w = coord[0][0], coord[1][0] if self.replace_sampling: self.map[coord_h, coord_w] = -1 return [coord_w, coord_h] class PainterOptimizer: def __init__(self, renderer: Painter, style: str, num_iter: int, lr_config: omegaconf.DictConfig, trainable_bg: bool = False): self.renderer = renderer self.num_iter = num_iter self.trainable_bg = trainable_bg self.lr_config = lr_config # set optimized params via style self.optim_point, self.optim_color, self.optim_width = { "iconography": (True, True, False), "pixelart": (False, True, False), "low-poly": (True, True, False), "sketch": (True, True, False), "ink": (True, True, True), "painting": (True, True, True) }.get(style, (False, False, False)) self.optim_bg = trainable_bg # set lr schedule schedule_cfg = lr_config.schedule if schedule_cfg.name == 'linear': self.lr_lambda = LinearDecayWithKeepLRLambda(init_lr=lr_config.point, keep_ratio=schedule_cfg.keep_ratio, decay_every=self.num_iter, decay_ratio=schedule_cfg.decay_ratio) elif schedule_cfg.name == 'cosine': self.lr_lambda = CosineWithWarmupLRLambda(num_steps=self.num_iter, warmup_steps=schedule_cfg.warmup_steps, warmup_start_lr=schedule_cfg.warmup_start_lr, warmup_end_lr=schedule_cfg.warmup_end_lr, cosine_end_lr=schedule_cfg.cosine_end_lr) else: print(f"{schedule_cfg.name} is not support.") self.lr_lambda = None self.point_optimizer = None self.color_optimizer = None self.width_optimizer = None self.bg_optimizer = None self.point_scheduler = None def init_optimizers(self, pid_delta: int = 0): # optimizer optim_cfg = self.lr_config.optim optim_name = optim_cfg.name params = {} if self.optim_point: self.renderer.set_point_parameters(pid_delta) params['point'] = self.renderer.get_point_parameters() self.point_optimizer = get_optimizer(optim_name, params['point'], self.lr_config.point, optim_cfg) if self.optim_color: self.renderer.set_color_parameters() params['color'] = self.renderer.get_color_parameters() self.color_optimizer = get_optimizer(optim_name, params['color'], self.lr_config.color, optim_cfg) if self.optim_width: self.renderer.set_width_parameters() params['width'] = self.renderer.get_width_parameters() if len(params['width']) > 0: self.width_optimizer = get_optimizer(optim_name, params['width'], self.lr_config.width, optim_cfg) if self.optim_bg: self.renderer.para_bg.requires_grad = True self.bg_optimizer = get_optimizer(optim_name, self.renderer.para_bg, self.lr_config.bg, optim_cfg) # lr schedule if self.lr_lambda is not None and self.optim_point: self.point_scheduler = LambdaLR(self.point_optimizer, lr_lambda=self.lr_lambda, last_epoch=-1) def update_lr(self): if self.point_scheduler is not None: self.point_scheduler.step() def zero_grad_(self): if self.point_optimizer is not None: self.point_optimizer.zero_grad() if self.color_optimizer is not None: self.color_optimizer.zero_grad() if self.width_optimizer is not None: self.width_optimizer.zero_grad() if self.bg_optimizer is not None: self.bg_optimizer.zero_grad() def step_(self): if self.point_optimizer is not None: self.point_optimizer.step() if self.color_optimizer is not None: self.color_optimizer.step() if self.width_optimizer is not None: self.width_optimizer.step() if self.bg_optimizer is not None: self.bg_optimizer.step() def get_lr(self) -> Dict: lr = {} if self.point_optimizer is not None: lr['pnt'] = self.point_optimizer.param_groups[0]['lr'] if self.color_optimizer is not None: lr['clr'] = self.color_optimizer.param_groups[0]['lr'] if self.width_optimizer is not None: lr['wd'] = self.width_optimizer.param_groups[0]['lr'] if self.bg_optimizer is not None: lr['bg'] = self.bg_optimizer.param_groups[0]['lr'] return lr class LinearDecayWithKeepLRLambda: """apply in LIVE stage""" def __init__(self, init_lr, keep_ratio, decay_every, decay_ratio): self.init_lr = init_lr self.keep_ratio = keep_ratio self.decay_every = decay_every self.decay_ratio = decay_ratio def __call__(self, n): if n < self.keep_ratio * self.decay_every: return self.init_lr decay_time = n // self.decay_every decay_step = n % self.decay_every lr_s = self.decay_ratio ** decay_time lr_e = self.decay_ratio ** (decay_time + 1) r = decay_step / self.decay_every lr = lr_s * (1 - r) + lr_e * r return lr class CosineWithWarmupLRLambda: """apply in fine-tuning stage""" def __init__(self, num_steps, warmup_steps, warmup_start_lr, warmup_end_lr, cosine_end_lr): self.n_steps = num_steps self.n_warmup = warmup_steps self.warmup_start_lr = warmup_start_lr self.warmup_end_lr = warmup_end_lr self.cosine_end_lr = cosine_end_lr def __call__(self, n): if n < self.n_warmup: # linearly warmup return self.warmup_start_lr + (n / self.n_warmup) * (self.warmup_end_lr - self.warmup_start_lr) else: # cosine decayed schedule return self.cosine_end_lr + 0.5 * (self.warmup_end_lr - self.cosine_end_lr) * ( 1 + math.cos(math.pi * (n - self.n_warmup) / (self.n_steps - self.n_warmup)))