741 lines
25 KiB
Python
741 lines
25 KiB
Python
# -*- coding: utf-8 -*-
|
||
# @Author: Peng Xiang
|
||
|
||
import types
|
||
import torch
|
||
import torch.nn.functional as F
|
||
import numpy as np
|
||
from torch import nn, einsum
|
||
from pointnet2_ops.pointnet2_utils import furthest_point_sample, \
|
||
gather_operation, ball_query, three_nn, three_interpolate, grouping_operation
|
||
|
||
class Conv1d(nn.Module):
|
||
def __init__(self, in_channel, out_channel, kernel_size=1, stride=1, if_bn=True, activation_fn=torch.relu):
|
||
super(Conv1d, self).__init__()
|
||
self.conv = nn.Conv1d(in_channel, out_channel, kernel_size, stride=stride)
|
||
self.if_bn = if_bn
|
||
self.bn = nn.BatchNorm1d(out_channel)
|
||
self.activation_fn = activation_fn
|
||
|
||
def forward(self, input):
|
||
out = self.conv(input)
|
||
if self.if_bn:
|
||
out = self.bn(out)
|
||
|
||
if self.activation_fn is not None:
|
||
out = self.activation_fn(out)
|
||
|
||
return out
|
||
|
||
class Conv2d(nn.Module):
|
||
def __init__(self, in_channel, out_channel, kernel_size=(1, 1), stride=(1, 1), if_bn=True, activation_fn=torch.relu):
|
||
super(Conv2d, self).__init__()
|
||
self.conv = nn.Conv2d(in_channel, out_channel, kernel_size, stride=stride)
|
||
self.if_bn = if_bn
|
||
self.bn = nn.BatchNorm2d(out_channel)
|
||
self.activation_fn = activation_fn
|
||
|
||
def forward(self, input):
|
||
out = self.conv(input)
|
||
if self.if_bn:
|
||
out = self.bn(out)
|
||
|
||
if self.activation_fn is not None:
|
||
out = self.activation_fn(out)
|
||
|
||
return out
|
||
|
||
class MLP(nn.Module):
|
||
def __init__(self, in_channel, layer_dims, bn=None):
|
||
super(MLP, self).__init__()
|
||
layers = []
|
||
last_channel = in_channel
|
||
for out_channel in layer_dims[:-1]:
|
||
layers.append(nn.Linear(last_channel, out_channel))
|
||
if bn:
|
||
layers.append(nn.BatchNorm1d(out_channel))
|
||
layers.append(nn.ReLU())
|
||
last_channel = out_channel
|
||
layers.append(nn.Linear(last_channel, layer_dims[-1]))
|
||
self.mlp = nn.Sequential(*layers)
|
||
|
||
def forward(self, inputs):
|
||
return self.mlp(inputs)
|
||
|
||
class MLP_CONV(nn.Module):
|
||
def __init__(self, in_channel, layer_dims, bn=None):
|
||
super(MLP_CONV, self).__init__()
|
||
layers = []
|
||
last_channel = in_channel
|
||
for out_channel in layer_dims[:-1]:
|
||
layers.append(nn.Conv1d(last_channel, out_channel, 1))
|
||
if bn:
|
||
layers.append(nn.BatchNorm1d(out_channel))
|
||
layers.append(nn.ReLU())
|
||
last_channel = out_channel
|
||
layers.append(nn.Conv1d(last_channel, layer_dims[-1], 1))
|
||
self.mlp = nn.Sequential(*layers)
|
||
|
||
def forward(self, inputs):
|
||
return self.mlp(inputs)
|
||
|
||
class MLP_Res(nn.Module):
|
||
def __init__(self, in_dim=128, hidden_dim=None, out_dim=128):
|
||
super(MLP_Res, self).__init__()
|
||
if hidden_dim is None:
|
||
hidden_dim = in_dim
|
||
self.conv_1 = nn.Conv1d(in_dim, hidden_dim, 1)
|
||
self.conv_2 = nn.Conv1d(hidden_dim, out_dim, 1)
|
||
self.conv_shortcut = nn.Conv1d(in_dim, out_dim, 1)
|
||
|
||
def forward(self, x):
|
||
"""
|
||
Args:
|
||
x: (B, out_dim, n)
|
||
"""
|
||
shortcut = self.conv_shortcut(x)
|
||
out = self.conv_2(torch.relu(self.conv_1(x))) + shortcut
|
||
return out
|
||
|
||
|
||
def sample_and_group(xyz, points, npoint, nsample, radius, use_xyz=True):
|
||
"""
|
||
Args:
|
||
xyz: Tensor, (B, 3, N)
|
||
points: Tensor, (B, f, N)
|
||
npoint: int
|
||
nsample: int
|
||
radius: float
|
||
use_xyz: boolean
|
||
|
||
Returns:
|
||
new_xyz: Tensor, (B, 3, npoint)
|
||
new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample)
|
||
idx_local: Tensor, (B, npoint, nsample)
|
||
grouped_xyz: Tensor, (B, 3, npoint, nsample)
|
||
|
||
"""
|
||
xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3)
|
||
new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint)
|
||
|
||
idx = ball_query(radius, nsample, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous()) # (B, npoint, nsample)
|
||
grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample)
|
||
grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, nsample)
|
||
|
||
if points is not None:
|
||
grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample)
|
||
if use_xyz:
|
||
new_points = torch.cat([grouped_xyz, grouped_points], 1)
|
||
else:
|
||
new_points = grouped_points
|
||
else:
|
||
new_points = grouped_xyz
|
||
|
||
return new_xyz, new_points, idx, grouped_xyz
|
||
|
||
|
||
def sample_and_group_all(xyz, points, use_xyz=True):
|
||
"""
|
||
Args:
|
||
xyz: Tensor, (B, 3, nsample)
|
||
points: Tensor, (B, f, nsample)
|
||
use_xyz: boolean
|
||
|
||
Returns:
|
||
new_xyz: Tensor, (B, 3, 1)
|
||
new_points: Tensor, (B, f|f+3|3, 1, nsample)
|
||
idx: Tensor, (B, 1, nsample)
|
||
grouped_xyz: Tensor, (B, 3, 1, nsample)
|
||
"""
|
||
b, _, nsample = xyz.shape
|
||
device = xyz.device
|
||
new_xyz = torch.zeros((1, 3, 1), dtype=torch.float, device=device).repeat(b, 1, 1)
|
||
grouped_xyz = xyz.reshape((b, 3, 1, nsample))
|
||
idx = torch.arange(nsample, device=device).reshape(1, 1, nsample).repeat(b, 1, 1)
|
||
if points is not None:
|
||
if use_xyz:
|
||
new_points = torch.cat([xyz, points], 1)
|
||
else:
|
||
new_points = points
|
||
new_points = new_points.unsqueeze(2)
|
||
else:
|
||
new_points = grouped_xyz
|
||
|
||
return new_xyz, new_points, idx, grouped_xyz
|
||
|
||
|
||
class PointNet_SA_Module(nn.Module):
|
||
def __init__(self, npoint, nsample, radius, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True):
|
||
"""
|
||
Args:
|
||
npoint: int, number of points to sample
|
||
nsample: int, number of points in each local region
|
||
radius: float
|
||
in_channel: int, input channel of features(points)
|
||
mlp: list of int,
|
||
"""
|
||
super(PointNet_SA_Module, self).__init__()
|
||
self.npoint = npoint
|
||
self.nsample = nsample
|
||
self.radius = radius
|
||
self.mlp = mlp
|
||
self.group_all = group_all
|
||
self.use_xyz = use_xyz
|
||
if use_xyz:
|
||
in_channel += 3
|
||
|
||
last_channel = in_channel
|
||
self.mlp_conv = []
|
||
for out_channel in mlp:
|
||
self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn))
|
||
last_channel = out_channel
|
||
|
||
self.mlp_conv = nn.Sequential(*self.mlp_conv)
|
||
|
||
def forward(self, xyz, points):
|
||
"""
|
||
Args:
|
||
xyz: Tensor, (B, 3, N)
|
||
points: Tensor, (B, f, N)
|
||
|
||
Returns:
|
||
new_xyz: Tensor, (B, 3, npoint)
|
||
new_points: Tensor, (B, mlp[-1], npoint)
|
||
"""
|
||
if self.group_all:
|
||
new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz)
|
||
else:
|
||
new_xyz, new_points, idx, grouped_xyz = sample_and_group(xyz, points, self.npoint, self.nsample, self.radius, self.use_xyz)
|
||
|
||
new_points = self.mlp_conv(new_points)
|
||
new_points = torch.max(new_points, 3)[0]
|
||
|
||
return new_xyz, new_points
|
||
|
||
|
||
class PointNet_FP_Module(nn.Module):
|
||
def __init__(self, in_channel, mlp, use_points1=False, in_channel_points1=None, if_bn=True):
|
||
"""
|
||
Args:
|
||
in_channel: int, input channel of points2
|
||
mlp: list of int
|
||
use_points1: boolean, if use points
|
||
in_channel_points1: int, input channel of points1
|
||
"""
|
||
super(PointNet_FP_Module, self).__init__()
|
||
self.use_points1 = use_points1
|
||
|
||
if use_points1:
|
||
in_channel += in_channel_points1
|
||
|
||
last_channel = in_channel
|
||
self.mlp_conv = []
|
||
for out_channel in mlp:
|
||
self.mlp_conv.append(Conv1d(last_channel, out_channel, if_bn=if_bn))
|
||
last_channel = out_channel
|
||
|
||
self.mlp_conv = nn.Sequential(*self.mlp_conv)
|
||
|
||
def forward(self, xyz1, xyz2, points1, points2):
|
||
"""
|
||
Args:
|
||
xyz1: Tensor, (B, 3, N)
|
||
xyz2: Tensor, (B, 3, M)
|
||
points1: Tensor, (B, in_channel, N)
|
||
points2: Tensor, (B, in_channel, M)
|
||
|
||
Returns:MLP_CONV
|
||
new_points: Tensor, (B, mlp[-1], N)
|
||
"""
|
||
dist, idx = three_nn(xyz1.permute(0, 2, 1).contiguous(), xyz2.permute(0, 2, 1).contiguous())
|
||
dist = torch.clamp_min(dist, 1e-10) # (B, N, 3)
|
||
recip_dist = 1.0/dist
|
||
norm = torch.sum(recip_dist, 2, keepdim=True).repeat((1, 1, 3))
|
||
weight = recip_dist / norm
|
||
interpolated_points = three_interpolate(points2, idx, weight) # B, in_channel, N
|
||
|
||
if self.use_points1:
|
||
new_points = torch.cat([interpolated_points, points1], 1)
|
||
else:
|
||
new_points = interpolated_points
|
||
|
||
new_points = self.mlp_conv(new_points)
|
||
return new_points
|
||
|
||
|
||
def square_distance(src, dst):
|
||
"""
|
||
Calculate Euclid distance between each two points.
|
||
|
||
src^T * dst = xn * xm + yn * ym + zn * zm；
|
||
sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
|
||
sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
|
||
dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
|
||
= sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
|
||
|
||
Input:
|
||
src: source points, [B, N, C]
|
||
dst: target points, [B, M, C]
|
||
Output:
|
||
dist: per-point square distance, [B, N, M]
|
||
"""
|
||
B, N, _ = src.shape
|
||
_, M, _ = dst.shape
|
||
dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) # B, N, M
|
||
dist += torch.sum(src ** 2, -1).view(B, N, 1)
|
||
dist += torch.sum(dst ** 2, -1).view(B, 1, M)
|
||
return dist
|
||
|
||
|
||
def query_knn(nsample, xyz, new_xyz, include_self=True):
|
||
"""Find k-NN of new_xyz in xyz"""
|
||
pad = 0 if include_self else 1
|
||
sqrdists = square_distance(new_xyz, xyz) # B, S, N
|
||
idx = torch.argsort(sqrdists, dim=-1, descending=False)[:, :, pad: nsample+pad]
|
||
return idx.int()
|
||
|
||
|
||
def sample_and_group_knn(xyz, points, npoint, k, use_xyz=True, idx=None):
|
||
"""
|
||
Args:
|
||
xyz: Tensor, (B, 3, N)
|
||
points: Tensor, (B, f, N)
|
||
npoint: int
|
||
nsample: int
|
||
radius: float
|
||
use_xyz: boolean
|
||
|
||
Returns:
|
||
new_xyz: Tensor, (B, 3, npoint)
|
||
new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample)
|
||
idx_local: Tensor, (B, npoint, nsample)
|
||
grouped_xyz: Tensor, (B, 3, npoint, nsample)
|
||
|
||
"""
|
||
xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3)
|
||
new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint)
|
||
if idx is None:
|
||
idx = query_knn(k, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous())
|
||
grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample)
|
||
grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, k)
|
||
|
||
if points is not None:
|
||
grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample)
|
||
if use_xyz:
|
||
new_points = torch.cat([grouped_xyz, grouped_points], 1)
|
||
else:
|
||
new_points = grouped_points
|
||
else:
|
||
new_points = grouped_xyz
|
||
|
||
return new_xyz, new_points, idx, grouped_xyz
|
||
|
||
|
||
class PointNet_SA_Module_KNN(nn.Module):
|
||
def __init__(self, npoint, nsample, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True, if_idx=False):
|
||
"""
|
||
Args:
|
||
npoint: int, number of points to sample
|
||
nsample: int, number of points in each local region
|
||
radius: float
|
||
in_channel: int, input channel of features(points)
|
||
mlp: list of int,
|
||
"""
|
||
super(PointNet_SA_Module_KNN, self).__init__()
|
||
self.npoint = npoint
|
||
self.nsample = nsample
|
||
self.mlp = mlp
|
||
self.group_all = group_all
|
||
self.use_xyz = use_xyz
|
||
self.if_idx = if_idx
|
||
if use_xyz:
|
||
in_channel += 3
|
||
|
||
last_channel = in_channel
|
||
self.mlp_conv = []
|
||
for out_channel in mlp[:-1]:
|
||
self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn))
|
||
last_channel = out_channel
|
||
self.mlp_conv.append(Conv2d(last_channel, mlp[-1], if_bn=False, activation_fn=None))
|
||
self.mlp_conv = nn.Sequential(*self.mlp_conv)
|
||
|
||
def forward(self, xyz, points, idx=None):
|
||
"""
|
||
Args:
|
||
xyz: Tensor, (B, 3, N)
|
||
points: Tensor, (B, f, N)
|
||
|
||
Returns:
|
||
new_xyz: Tensor, (B, 3, npoint)
|
||
new_points: Tensor, (B, mlp[-1], npoint)
|
||
"""
|
||
if self.group_all:
|
||
new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz)
|
||
else:
|
||
new_xyz, new_points, idx, grouped_xyz = sample_and_group_knn(xyz, points, self.npoint, self.nsample, self.use_xyz, idx=idx)
|
||
|
||
new_points = self.mlp_conv(new_points)
|
||
new_points = torch.max(new_points, 3)[0]
|
||
|
||
if self.if_idx:
|
||
return new_xyz, new_points, idx
|
||
else:
|
||
return new_xyz, new_points
|
||
|
||
|
||
def fps_subsample(pcd, n_points=2048):
|
||
"""
|
||
Args
|
||
pcd: (b, 16384, 3)
|
||
|
||
returns
|
||
new_pcd: (b, n_points, 3)
|
||
"""
|
||
new_pcd = gather_operation(pcd.permute(0, 2, 1).contiguous(), furthest_point_sample(pcd, n_points))
|
||
new_pcd = new_pcd.permute(0, 2, 1).contiguous()
|
||
return new_pcd
|
||
|
||
|
||
class Transformer(nn.Module):
|
||
def __init__(self, in_channel, dim=256, n_knn=16, pos_hidden_dim=64, attn_hidden_multiplier=4):
|
||
super(Transformer, self).__init__()
|
||
self.n_knn = n_knn
|
||
self.conv_key = nn.Conv1d(dim, dim, 1)
|
||
self.conv_query = nn.Conv1d(dim, dim, 1)
|
||
self.conv_value = nn.Conv1d(dim, dim, 1)
|
||
|
||
self.pos_mlp = nn.Sequential(
|
||
nn.Conv2d(3, pos_hidden_dim, 1),
|
||
nn.BatchNorm2d(pos_hidden_dim),
|
||
nn.ReLU(),
|
||
nn.Conv2d(pos_hidden_dim, dim, 1)
|
||
)
|
||
|
||
self.attn_mlp = nn.Sequential(
|
||
nn.Conv2d(dim, dim * attn_hidden_multiplier, 1),
|
||
nn.BatchNorm2d(dim * attn_hidden_multiplier),
|
||
nn.ReLU(),
|
||
nn.Conv2d(dim * attn_hidden_multiplier, dim, 1)
|
||
)
|
||
|
||
self.linear_start = nn.Conv1d(in_channel, dim, 1)
|
||
self.linear_end = nn.Conv1d(dim, in_channel, 1)
|
||
|
||
def forward(self, x, pos):
|
||
"""feed forward of transformer
|
||
Args:
|
||
x: Tensor of features, (B, in_channel, n)
|
||
pos: Tensor of positions, (B, 3, n)
|
||
|
||
Returns:
|
||
y: Tensor of features with attention, (B, in_channel, n)
|
||
"""
|
||
|
||
identity = x
|
||
|
||
x = self.linear_start(x)
|
||
b, dim, n = x.shape
|
||
|
||
pos_flipped = pos.permute(0, 2, 1).contiguous()
|
||
idx_knn = query_knn(self.n_knn, pos_flipped, pos_flipped)
|
||
key = self.conv_key(x)
|
||
value = self.conv_value(x)
|
||
query = self.conv_query(x)
|
||
|
||
key = grouping_operation(key, idx_knn) # b, dim, n, n_knn
|
||
qk_rel = query.reshape((b, -1, n, 1)) - key
|
||
|
||
pos_rel = pos.reshape((b, -1, n, 1)) - grouping_operation(pos, idx_knn) # b, 3, n, n_knn
|
||
pos_embedding = self.pos_mlp(pos_rel) # b, dim, n, n_knn
|
||
|
||
attention = self.attn_mlp(qk_rel + pos_embedding)
|
||
attention = torch.softmax(attention, -1)
|
||
|
||
value = value.reshape((b, -1, n, 1)) + pos_embedding
|
||
|
||
agg = einsum('b c i j, b c i j -> b c i', attention, value) # b, dim, n
|
||
y = self.linear_end(agg)
|
||
|
||
return y+identity
|
||
|
||
|
||
class CouplingLayer(nn.Module):
|
||
|
||
def __init__(self, d, intermediate_dim, swap=False):
|
||
nn.Module.__init__(self)
|
||
self.d = d - (d // 2)
|
||
self.swap = swap
|
||
self.net_s_t = nn.Sequential(
|
||
nn.Linear(self.d, intermediate_dim),
|
||
nn.ReLU(inplace=True),
|
||
nn.Linear(intermediate_dim, intermediate_dim),
|
||
nn.ReLU(inplace=True),
|
||
nn.Linear(intermediate_dim, (d - self.d) * 2),
|
||
)
|
||
|
||
def forward(self, x, logpx=None, reverse=False):
|
||
|
||
if self.swap:
|
||
x = torch.cat([x[:, self.d:], x[:, :self.d]], 1)
|
||
|
||
in_dim = self.d
|
||
out_dim = x.shape[1] - self.d
|
||
|
||
s_t = self.net_s_t(x[:, :in_dim])
|
||
scale = torch.sigmoid(s_t[:, :out_dim] + 2.)
|
||
shift = s_t[:, out_dim:]
|
||
|
||
logdetjac = torch.sum(torch.log(scale).view(scale.shape[0], -1), 1, keepdim=True)
|
||
|
||
if not reverse:
|
||
y1 = x[:, self.d:] * scale + shift
|
||
delta_logp = -logdetjac
|
||
else:
|
||
y1 = (x[:, self.d:] - shift) / scale
|
||
delta_logp = logdetjac
|
||
|
||
y = torch.cat([x[:, :self.d], y1], 1) if not self.swap else torch.cat([y1, x[:, :self.d]], 1)
|
||
|
||
if logpx is None:
|
||
return y
|
||
else:
|
||
return y, logpx + delta_logp
|
||
|
||
|
||
class SequentialFlow(nn.Module):
|
||
"""A generalized nn.Sequential container for normalizing flows.
|
||
"""
|
||
|
||
def __init__(self, layersList):
|
||
super(SequentialFlow, self).__init__()
|
||
self.chain = nn.ModuleList(layersList)
|
||
|
||
def forward(self, x, logpx=None, reverse=False, inds=None):
|
||
if inds is None:
|
||
if reverse:
|
||
inds = range(len(self.chain) - 1, -1, -1)
|
||
else:
|
||
inds = range(len(self.chain))
|
||
|
||
if logpx is None:
|
||
for i in inds:
|
||
x = self.chain[i](x, reverse=reverse)
|
||
return x
|
||
else:
|
||
for i in inds:
|
||
x, logpx = self.chain[i](x, logpx, reverse=reverse)
|
||
return x, logpx
|
||
|
||
|
||
def build_latent_flow(args):
|
||
chain = []
|
||
for i in range(args.latent_flow_depth):
|
||
chain.append(CouplingLayer(args.latent_dim, args.latent_flow_hidden_dim, swap=(i % 2 == 0)))
|
||
return SequentialFlow(chain)
|
||
|
||
|
||
##################
|
||
## SpectralNorm ##
|
||
##################
|
||
|
||
POWER_ITERATION_FN = "spectral_norm_power_iteration"
|
||
|
||
|
||
class SpectralNorm(object):
|
||
def __init__(self, name='weight', dim=0, eps=1e-12):
|
||
self.name = name
|
||
self.dim = dim
|
||
self.eps = eps
|
||
|
||
def compute_weight(self, module, n_power_iterations):
|
||
if n_power_iterations < 0:
|
||
raise ValueError(
|
||
'Expected n_power_iterations to be non-negative, but '
|
||
'got n_power_iterations={}'.format(n_power_iterations)
|
||
)
|
||
|
||
weight = getattr(module, self.name + '_orig')
|
||
u = getattr(module, self.name + '_u')
|
||
v = getattr(module, self.name + '_v')
|
||
weight_mat = weight
|
||
if self.dim != 0:
|
||
# permute dim to front
|
||
weight_mat = weight_mat.permute(self.dim, * [d for d in range(weight_mat.dim()) if d != self.dim])
|
||
height = weight_mat.size(0)
|
||
weight_mat = weight_mat.reshape(height, -1)
|
||
with torch.no_grad():
|
||
for _ in range(n_power_iterations):
|
||
# Spectral norm of weight equals to `u^T W v`, where `u` and `v`
|
||
# are the first left and right singular vectors.
|
||
# This power iteration produces approximations of `u` and `v`.
|
||
v = F.normalize(torch.matmul(weight_mat.t(), u), dim=0, eps=self.eps)
|
||
u = F.normalize(torch.matmul(weight_mat, v), dim=0, eps=self.eps)
|
||
setattr(module, self.name + '_u', u)
|
||
setattr(module, self.name + '_v', v)
|
||
|
||
sigma = torch.dot(u, torch.matmul(weight_mat, v))
|
||
weight = weight / sigma
|
||
setattr(module, self.name, weight)
|
||
|
||
def remove(self, module):
|
||
weight = getattr(module, self.name)
|
||
delattr(module, self.name)
|
||
delattr(module, self.name + '_u')
|
||
delattr(module, self.name + '_orig')
|
||
module.register_parameter(self.name, torch.nn.Parameter(weight))
|
||
|
||
def get_update_method(self, module):
|
||
def update_fn(module, n_power_iterations):
|
||
self.compute_weight(module, n_power_iterations)
|
||
|
||
return update_fn
|
||
|
||
def __call__(self, module, unused_inputs):
|
||
del unused_inputs
|
||
self.compute_weight(module, n_power_iterations=0)
|
||
|
||
# requires_grad might be either True or False during inference.
|
||
if not module.training:
|
||
r_g = getattr(module, self.name + '_orig').requires_grad
|
||
setattr(module, self.name, getattr(module, self.name).detach().requires_grad_(r_g))
|
||
|
||
@staticmethod
|
||
def apply(module, name, dim, eps):
|
||
fn = SpectralNorm(name, dim, eps)
|
||
weight = module._parameters[name]
|
||
height = weight.size(dim)
|
||
|
||
u = F.normalize(weight.new_empty(height).normal_(0, 1), dim=0, eps=fn.eps)
|
||
v = F.normalize(weight.new_empty(int(weight.numel() / height)).normal_(0, 1), dim=0, eps=fn.eps)
|
||
delattr(module, fn.name)
|
||
module.register_parameter(fn.name + "_orig", weight)
|
||
# We still need to assign weight back as fn.name because all sorts of
|
||
# things may assume that it exists, e.g., when initializing weights.
|
||
# However, we can't directly assign as it could be an nn.Parameter and
|
||
# gets added as a parameter. Instead, we register weight.data as a
|
||
# buffer, which will cause weight to be included in the state dict
|
||
# and also supports nn.init due to shared storage.
|
||
module.register_buffer(fn.name, weight.data)
|
||
module.register_buffer(fn.name + "_u", u)
|
||
module.register_buffer(fn.name + "_v", v)
|
||
|
||
setattr(module, POWER_ITERATION_FN, types.MethodType(fn.get_update_method(module), module))
|
||
|
||
module.register_forward_pre_hook(fn)
|
||
return fn
|
||
|
||
|
||
def inplace_spectral_norm(module, name='weight', dim=None, eps=1e-12):
|
||
r"""Applies spectral normalization to a parameter in the given module.
|
||
.. math::
|
||
\mathbf{W} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})} \\
|
||
\sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2}
|
||
Spectral normalization stabilizes the training of discriminators (critics)
|
||
in Generaive Adversarial Networks (GANs) by rescaling the weight tensor
|
||
with spectral norm :math:`\sigma` of the weight matrix calculated using
|
||
power iteration method. If the dimension of the weight tensor is greater
|
||
than 2, it is reshaped to 2D in power iteration method to get spectral
|
||
norm. This is implemented via a hook that calculates spectral norm and
|
||
rescales weight before every :meth:`~Module.forward` call.
|
||
See `Spectral Normalization for Generative Adversarial Networks`_ .
|
||
.. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957
|
||
Args:
|
||
module (nn.Module): containing module
|
||
name (str, optional): name of weight parameter
|
||
n_power_iterations (int, optional): number of power iterations to
|
||
calculate spectal norm
|
||
dim (int, optional): dimension corresponding to number of outputs,
|
||
the default is 0, except for modules that are instances of
|
||
ConvTranspose1/2/3d, when it is 1
|
||
eps (float, optional): epsilon for numerical stability in
|
||
calculating norms
|
||
Returns:
|
||
The original module with the spectal norm hook
|
||
Example::
|
||
>>> m = spectral_norm(nn.Linear(20, 40))
|
||
Linear (20 -> 40)
|
||
>>> m.weight_u.size()
|
||
torch.Size([20])
|
||
"""
|
||
if dim is None:
|
||
if isinstance(module, (torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d)):
|
||
dim = 1
|
||
else:
|
||
dim = 0
|
||
SpectralNorm.apply(module, name, dim=dim, eps=eps)
|
||
return module
|
||
|
||
|
||
def remove_spectral_norm(module, name='weight'):
|
||
r"""Removes the spectral normalization reparameterization from a module.
|
||
Args:
|
||
module (nn.Module): containing module
|
||
name (str, optional): name of weight parameter
|
||
Example:
|
||
>>> m = spectral_norm(nn.Linear(40, 10))
|
||
>>> remove_spectral_norm(m)
|
||
"""
|
||
for k, hook in module._forward_pre_hooks.items():
|
||
if isinstance(hook, SpectralNorm) and hook.name == name:
|
||
hook.remove(module)
|
||
del module._forward_pre_hooks[k]
|
||
return module
|
||
|
||
raise ValueError("spectral_norm of '{}' not found in {}".format(name, module))
|
||
|
||
|
||
def add_spectral_norm(model, logger=None):
|
||
"""Applies spectral norm to all modules within the scope of a CNF."""
|
||
|
||
def apply_spectral_norm(module):
|
||
if 'weight' in module._parameters:
|
||
if logger: logger.info("Adding spectral norm to {}".format(module))
|
||
inplace_spectral_norm(module, 'weight')
|
||
|
||
def find_coupling_layer(module):
|
||
if isinstance(module, CouplingLayer):
|
||
module.apply(apply_spectral_norm)
|
||
else:
|
||
for child in module.children():
|
||
find_coupling_layer(child)
|
||
|
||
find_coupling_layer(model)
|
||
|
||
|
||
def spectral_norm_power_iteration(model, n_power_iterations=1):
|
||
|
||
def recursive_power_iteration(module):
|
||
if hasattr(module, POWER_ITERATION_FN):
|
||
getattr(module, POWER_ITERATION_FN)(n_power_iterations)
|
||
|
||
model.apply(recursive_power_iteration)
|
||
|
||
def reparameterize_gaussian(mean, logvar):
|
||
std = torch.exp(0.5 * logvar)
|
||
eps = torch.randn(std.size()).to(mean)
|
||
return mean + std * eps
|
||
|
||
|
||
|
||
def gaussian_entropy(logvar):
|
||
const = 0.5 * float(logvar.size(1)) * (1. + np.log(np.pi * 2))
|
||
ent = 0.5 * logvar.sum(dim=1, keepdim=False) + const
|
||
return ent
|
||
|
||
|
||
def standard_normal_logprob(z):
|
||
dim = z.size(-1)
|
||
log_z = -0.5 * dim * np.log(2 * np.pi)
|
||
return log_z - z.pow(2) / 2
|
||
|
||
def truncated_normal_(tensor, mean=0, std=1, trunc_std=2):
|
||
"""
|
||
Taken from https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15
|
||
"""
|
||
size = tensor.shape
|
||
tmp = tensor.new_empty(size + (4,)).normal_()
|
||
valid = (tmp < trunc_std) & (tmp > -trunc_std)
|
||
ind = valid.max(-1, keepdim=True)[1]
|
||
tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
|
||
tensor.data.mul_(std).add_(mean)
|
||
return tensor |