PKUWilliamYang's picture
V1
983684c
raw
history blame
14.2 kB
import math
import torch
from torch import autograd
from torch.nn import functional as F
import numpy as np
from model.stylegan.distributed import reduce_sum
from model.stylegan.op import upfirdn2d
class AdaptiveAugment:
def __init__(self, ada_aug_target, ada_aug_len, update_every, device):
self.ada_aug_target = ada_aug_target
self.ada_aug_len = ada_aug_len
self.update_every = update_every
self.ada_update = 0
self.ada_aug_buf = torch.tensor([0.0, 0.0], device=device)
self.r_t_stat = 0
self.ada_aug_p = 0
@torch.no_grad()
def tune(self, real_pred):
self.ada_aug_buf += torch.tensor(
(torch.sign(real_pred).sum().item(), real_pred.shape[0]),
device=real_pred.device,
)
self.ada_update += 1
if self.ada_update % self.update_every == 0:
self.ada_aug_buf = reduce_sum(self.ada_aug_buf)
pred_signs, n_pred = self.ada_aug_buf.tolist()
self.r_t_stat = pred_signs / n_pred
if self.r_t_stat > self.ada_aug_target:
sign = 1
else:
sign = -1
self.ada_aug_p += sign * n_pred / self.ada_aug_len
self.ada_aug_p = min(1, max(0, self.ada_aug_p))
self.ada_aug_buf.mul_(0)
self.ada_update = 0
return self.ada_aug_p
SYM6 = (
0.015404109327027373,
0.0034907120842174702,
-0.11799011114819057,
-0.048311742585633,
0.4910559419267466,
0.787641141030194,
0.3379294217276218,
-0.07263752278646252,
-0.021060292512300564,
0.04472490177066578,
0.0017677118642428036,
-0.007800708325034148,
)
def translate_mat(t_x, t_y, device="cpu"):
batch = t_x.shape[0]
mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
translate = torch.stack((t_x, t_y), 1)
mat[:, :2, 2] = translate
return mat
def rotate_mat(theta, device="cpu"):
batch = theta.shape[0]
mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
sin_t = torch.sin(theta)
cos_t = torch.cos(theta)
rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2)
mat[:, :2, :2] = rot
return mat
def scale_mat(s_x, s_y, device="cpu"):
batch = s_x.shape[0]
mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
mat[:, 0, 0] = s_x
mat[:, 1, 1] = s_y
return mat
def translate3d_mat(t_x, t_y, t_z):
batch = t_x.shape[0]
mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
translate = torch.stack((t_x, t_y, t_z), 1)
mat[:, :3, 3] = translate
return mat
def rotate3d_mat(axis, theta):
batch = theta.shape[0]
u_x, u_y, u_z = axis
eye = torch.eye(3).unsqueeze(0)
cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0)
outer = torch.tensor(axis)
outer = (outer.unsqueeze(1) * outer).unsqueeze(0)
sin_t = torch.sin(theta).view(-1, 1, 1)
cos_t = torch.cos(theta).view(-1, 1, 1)
rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer
eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
eye_4[:, :3, :3] = rot
return eye_4
def scale3d_mat(s_x, s_y, s_z):
batch = s_x.shape[0]
mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
mat[:, 0, 0] = s_x
mat[:, 1, 1] = s_y
mat[:, 2, 2] = s_z
return mat
def luma_flip_mat(axis, i):
batch = i.shape[0]
eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
axis = torch.tensor(axis + (0,))
flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1)
return eye - flip
def saturation_mat(axis, i):
batch = i.shape[0]
eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
axis = torch.tensor(axis + (0,))
axis = torch.ger(axis, axis)
saturate = axis + (eye - axis) * i.view(-1, 1, 1)
return saturate
def lognormal_sample(size, mean=0, std=1, device="cpu"):
return torch.empty(size, device=device).log_normal_(mean=mean, std=std)
def category_sample(size, categories, device="cpu"):
category = torch.tensor(categories, device=device)
sample = torch.randint(high=len(categories), size=(size,), device=device)
return category[sample]
def uniform_sample(size, low, high, device="cpu"):
return torch.empty(size, device=device).uniform_(low, high)
def normal_sample(size, mean=0, std=1, device="cpu"):
return torch.empty(size, device=device).normal_(mean, std)
def bernoulli_sample(size, p, device="cpu"):
return torch.empty(size, device=device).bernoulli_(p)
def random_mat_apply(p, transform, prev, eye, device="cpu"):
size = transform.shape[0]
select = bernoulli_sample(size, p, device=device).view(size, 1, 1)
select_transform = select * transform + (1 - select) * eye
return select_transform @ prev
def sample_affine(p, size, height, width, device="cpu"):
G = torch.eye(3, device=device).unsqueeze(0).repeat(size, 1, 1)
eye = G
# flip
param = category_sample(size, (0, 1))
Gc = scale_mat(1 - 2.0 * param, torch.ones(size), device=device)
G = random_mat_apply(p, Gc, G, eye, device=device)
# print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\n')
# 90 rotate
#param = category_sample(size, (0, 3))
#Gc = rotate_mat(-math.pi / 2 * param, device=device)
#G = random_mat_apply(p, Gc, G, eye, device=device)
# print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\n')
# integer translate
param = uniform_sample(size, -0.125, 0.125)
param_height = torch.round(param * height) / height
param_width = torch.round(param * width) / width
Gc = translate_mat(param_width, param_height, device=device)
G = random_mat_apply(p, Gc, G, eye, device=device)
# print('integer translate', G, translate_mat(param_width, param_height), sep='\n')
# isotropic scale
param = lognormal_sample(size, std=0.2 * math.log(2))
Gc = scale_mat(param, param, device=device)
G = random_mat_apply(p, Gc, G, eye, device=device)
# print('isotropic scale', G, scale_mat(param, param), sep='\n')
p_rot = 1 - math.sqrt(1 - p)
# pre-rotate
param = uniform_sample(size, -math.pi, math.pi)
Gc = rotate_mat(-param, device=device)
G = random_mat_apply(p_rot, Gc, G, eye, device=device)
# print('pre-rotate', G, rotate_mat(-param), sep='\n')
# anisotropic scale
param = lognormal_sample(size, std=0.2 * math.log(2))
Gc = scale_mat(param, 1 / param, device=device)
G = random_mat_apply(p, Gc, G, eye, device=device)
# print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\n')
# post-rotate
param = uniform_sample(size, -math.pi, math.pi)
Gc = rotate_mat(-param, device=device)
G = random_mat_apply(p_rot, Gc, G, eye, device=device)
# print('post-rotate', G, rotate_mat(-param), sep='\n')
# fractional translate
param = normal_sample(size, std=0.125)
Gc = translate_mat(param, param, device=device)
G = random_mat_apply(p, Gc, G, eye, device=device)
# print('fractional translate', G, translate_mat(param, param), sep='\n')
return G
def sample_color(p, size):
C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1)
eye = C
axis_val = 1 / math.sqrt(3)
axis = (axis_val, axis_val, axis_val)
# brightness
param = normal_sample(size, std=0.2)
Cc = translate3d_mat(param, param, param)
C = random_mat_apply(p, Cc, C, eye)
# contrast
param = lognormal_sample(size, std=0.5 * math.log(2))
Cc = scale3d_mat(param, param, param)
C = random_mat_apply(p, Cc, C, eye)
# luma flip
param = category_sample(size, (0, 1))
Cc = luma_flip_mat(axis, param)
C = random_mat_apply(p, Cc, C, eye)
# hue rotation
param = uniform_sample(size, -math.pi, math.pi)
Cc = rotate3d_mat(axis, param)
C = random_mat_apply(p, Cc, C, eye)
# saturation
param = lognormal_sample(size, std=1 * math.log(2))
Cc = saturation_mat(axis, param)
C = random_mat_apply(p, Cc, C, eye)
return C
def make_grid(shape, x0, x1, y0, y1, device):
n, c, h, w = shape
grid = torch.empty(n, h, w, 3, device=device)
grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device)
grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1)
grid[:, :, :, 2] = 1
return grid
def affine_grid(grid, mat):
n, h, w, _ = grid.shape
return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2)
def get_padding(G, height, width, kernel_size):
device = G.device
cx = (width - 1) / 2
cy = (height - 1) / 2
cp = torch.tensor(
[(-cx, -cy, 1), (cx, -cy, 1), (cx, cy, 1), (-cx, cy, 1)], device=device
)
cp = G @ cp.T
pad_k = kernel_size // 4
pad = cp[:, :2, :].permute(1, 0, 2).flatten(1)
pad = torch.cat((-pad, pad)).max(1).values
pad = pad + torch.tensor([pad_k * 2 - cx, pad_k * 2 - cy] * 2, device=device)
pad = pad.max(torch.tensor([0, 0] * 2, device=device))
pad = pad.min(torch.tensor([width - 1, height - 1] * 2, device=device))
pad_x1, pad_y1, pad_x2, pad_y2 = pad.ceil().to(torch.int32)
return pad_x1, pad_x2, pad_y1, pad_y2
def try_sample_affine_and_pad(img, p, kernel_size, G=None):
batch, _, height, width = img.shape
G_try = G
if G is None:
G_try = torch.inverse(sample_affine(p, batch, height, width))
pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(G_try, height, width, kernel_size)
img_pad = F.pad(img, (pad_x1, pad_x2, pad_y1, pad_y2), mode="reflect")
return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2)
class GridSampleForward(autograd.Function):
@staticmethod
def forward(ctx, input, grid):
out = F.grid_sample(
input, grid, mode="bilinear", padding_mode="zeros", align_corners=False
)
ctx.save_for_backward(input, grid)
return out
@staticmethod
def backward(ctx, grad_output):
input, grid = ctx.saved_tensors
grad_input, grad_grid = GridSampleBackward.apply(grad_output, input, grid)
return grad_input, grad_grid
class GridSampleBackward(autograd.Function):
@staticmethod
def forward(ctx, grad_output, input, grid):
op = torch._C._jit_get_operation("aten::grid_sampler_2d_backward")
grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)
ctx.save_for_backward(grid)
return grad_input, grad_grid
@staticmethod
def backward(ctx, grad_grad_input, grad_grad_grid):
grid, = ctx.saved_tensors
grad_grad_output = None
if ctx.needs_input_grad[0]:
grad_grad_output = GridSampleForward.apply(grad_grad_input, grid)
return grad_grad_output, None, None
grid_sample = GridSampleForward.apply
def scale_mat_single(s_x, s_y):
return torch.tensor(((s_x, 0, 0), (0, s_y, 0), (0, 0, 1)), dtype=torch.float32)
def translate_mat_single(t_x, t_y):
return torch.tensor(((1, 0, t_x), (0, 1, t_y), (0, 0, 1)), dtype=torch.float32)
def random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6):
kernel = antialiasing_kernel
len_k = len(kernel)
kernel = torch.as_tensor(kernel).to(img)
# kernel = torch.ger(kernel, kernel).to(img)
kernel_flip = torch.flip(kernel, (0,))
img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad(
img, p, len_k, G
)
G_inv = (
translate_mat_single((pad_x1 - pad_x2).item() / 2, (pad_y1 - pad_y2).item() / 2)
@ G
)
up_pad = (
(len_k + 2 - 1) // 2,
(len_k - 2) // 2,
(len_k + 2 - 1) // 2,
(len_k - 2) // 2,
)
img_2x = upfirdn2d(img_pad, kernel.unsqueeze(0), up=(2, 1), pad=(*up_pad[:2], 0, 0))
img_2x = upfirdn2d(img_2x, kernel.unsqueeze(1), up=(1, 2), pad=(0, 0, *up_pad[2:]))
G_inv = scale_mat_single(2, 2) @ G_inv @ scale_mat_single(1 / 2, 1 / 2)
G_inv = translate_mat_single(-0.5, -0.5) @ G_inv @ translate_mat_single(0.5, 0.5)
batch_size, channel, height, width = img.shape
pad_k = len_k // 4
shape = (batch_size, channel, (height + pad_k * 2) * 2, (width + pad_k * 2) * 2)
G_inv = (
scale_mat_single(2 / img_2x.shape[3], 2 / img_2x.shape[2])
@ G_inv
@ scale_mat_single(1 / (2 / shape[3]), 1 / (2 / shape[2]))
)
grid = F.affine_grid(G_inv[:, :2, :].to(img_2x), shape, align_corners=False)
img_affine = grid_sample(img_2x, grid)
d_p = -pad_k * 2
down_pad = (
d_p + (len_k - 2 + 1) // 2,
d_p + (len_k - 2) // 2,
d_p + (len_k - 2 + 1) // 2,
d_p + (len_k - 2) // 2,
)
img_down = upfirdn2d(
img_affine, kernel_flip.unsqueeze(0), down=(2, 1), pad=(*down_pad[:2], 0, 0)
)
img_down = upfirdn2d(
img_down, kernel_flip.unsqueeze(1), down=(1, 2), pad=(0, 0, *down_pad[2:])
)
return img_down, G
def apply_color(img, mat):
batch = img.shape[0]
img = img.permute(0, 2, 3, 1)
mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3)
mat_add = mat[:, :3, 3].view(batch, 1, 1, 3)
img = img @ mat_mul + mat_add
img = img.permute(0, 3, 1, 2)
return img
def random_apply_color(img, p, C=None):
if C is None:
C = sample_color(p, img.shape[0])
img = apply_color(img, C.to(img))
return img, C
def augment(img, p, transform_matrix=(None, None)):
img, G = random_apply_affine(img, p, transform_matrix[0])
if img.shape[1] == 3:
img, C = random_apply_color(img, p, transform_matrix[1])
else:
tmp, C = random_apply_color(img[:,0:3], p, transform_matrix[1])
img = torch.cat((tmp, img[:,3:]), dim=1)
return img, (G, C)