|
import os |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from torch.utils.tensorboard import SummaryWriter |
|
|
|
|
|
class DNO(object): |
|
def __init__( |
|
self, |
|
optimize: bool, |
|
max_train_steps: int, |
|
learning_rate: float, |
|
lr_scheduler: str, |
|
lr_warmup_steps: int, |
|
clip_grad: bool, |
|
loss_hint_type: str, |
|
loss_diff_penalty: float, |
|
loss_correlate_penalty: float, |
|
visualize_samples: int, |
|
visualize_ske_steps: list[int], |
|
output_dir: str |
|
) -> None: |
|
|
|
self.optimize = optimize |
|
self.max_train_steps = max_train_steps |
|
self.learning_rate = learning_rate |
|
self.lr_scheduler = lr_scheduler |
|
self.lr_warmup_steps = lr_warmup_steps |
|
self.clip_grad = clip_grad |
|
self.loss_hint_type = loss_hint_type |
|
self.loss_diff_penalty = loss_diff_penalty |
|
self.loss_correlate_penalty = loss_correlate_penalty |
|
|
|
if loss_hint_type == 'l1': |
|
self.loss_hint_func = F.l1_loss |
|
elif loss_hint_type == 'l1_smooth': |
|
self.loss_hint_func = F.smooth_l1_loss |
|
elif loss_hint_type == 'l2': |
|
self.loss_hint_func = F.mse_loss |
|
else: |
|
raise ValueError(f'Invalid loss type: {loss_hint_type}') |
|
|
|
self.visualize_samples = float('inf') if visualize_samples == 'inf' else visualize_samples |
|
assert self.visualize_samples >= 0 |
|
self.visualize_samples_done = 0 |
|
self.visualize_ske_steps = visualize_ske_steps |
|
if len(visualize_ske_steps) > 0: |
|
self.vis_dir = os.path.join(output_dir, 'vis_optimize') |
|
os.makedirs(self.vis_dir) |
|
|
|
self.writer = None |
|
self.output_dir = output_dir |
|
if self.visualize_samples > 0: |
|
self.writer = SummaryWriter(output_dir) |
|
|
|
@property |
|
def do_visualize(self): |
|
return self.visualize_samples_done < self.visualize_samples |
|
|
|
@staticmethod |
|
def noise_regularize_1d(noise: torch.Tensor, stop_at: int = 2, dim: int = 1) -> torch.Tensor: |
|
size = noise.shape[dim] |
|
if size & (size - 1) != 0: |
|
new_size = 2 ** (size - 1).bit_length() |
|
pad = new_size - size |
|
pad_shape = list(noise.shape) |
|
pad_shape[dim] = pad |
|
pad_noise = torch.randn(*pad_shape, device=noise.device) |
|
noise = torch.cat([noise, pad_noise], dim=dim) |
|
size = noise.shape[dim] |
|
|
|
loss = torch.zeros(noise.shape[0], device=noise.device) |
|
while size > stop_at: |
|
rolled_noise = torch.roll(noise, shifts=1, dims=dim) |
|
loss += (noise * rolled_noise).mean(dim=tuple(range(1, noise.ndim))).pow(2) |
|
noise = noise.view(*noise.shape[:dim], size // 2, 2, *noise.shape[dim + 1:]).mean(dim=dim + 1) |
|
size //= 2 |
|
return loss |
|
|
