RenderNet模型复现
使用pytorch进行RenderNet模型的复现
没有使用z特征向量,z特征向量的解释可以参考论文Category Level Object Pose Estimation via Neural Analysis-by-Synthesis了解其具体含义,这里被定义为物体的外观特征。模型直接输入物体的体素,通过位姿控制生成对应的轮廓,其中数据集采用blender生成的轮廓。
import numpy as np
import torch
from torch.nn import init
import torch.utils.data
import torch.utils.data.distributed
from torch import nn
from torch.nn import functional as F
from torch.optim import lr_scheduler
from model.networks import spectral_norm
class Generator(nn.Module):
def __init__(self, z_dim, euler_seq='zyx', **kwargs):
super().__init__()
self.shape_code = nn.Parameter(0.02*torch.randn(1,512,4,4,4),requires_grad=True)
# Upsampling 3D
self.enc_1 = nn.Sequential(*[nn.ConvTranspose3d(512,128,kernel_size=4,stride=2, padding=1)])
self.enc_2 = nn.Sequential(*[nn.ConvTranspose3d(128,64, kernel_size=4, stride=2, padding=1)])
# Projection
self.proj = nn.Sequential(*[nn.ConvTranspose2d(64*16,64*16, kernel_size=1,stride=1)])
# Upsampling 2D
self.enc_3 = nn.Sequential(*[nn.ConvTranspose2d(64*16,64*4,kernel_size=4,stride=2, padding=1)])
self.enc_4 = nn.Sequential(*[nn.ConvTranspose2d(64*4,64,kernel_size=4,stride=2, padding=1)])
self.enc_5 = nn.Sequential(*[nn.ConvTranspose2d(64,3,kernel_size=3,stride=1,padding=1)])
# MLP for AdaIN
self.mlp0 = LinearBlock(z_dim,512*2,activation='relu')
self.mlp1 = LinearBlock(z_dim,128*2,activation='relu')
self.mlp2 = LinearBlock(z_dim,64*2,activation='relu')
self.mlp3 = LinearBlock(z_dim,256*2,activation='relu')
self.mlp4 = LinearBlock(z_dim,64*2,activation='relu')
self.euler_seq = euler_seq
def forward(self, z, angle, a=None, debug=False):
b,_ = z.size()
angle = angle / 180. * np.pi
# Upsampling 3D
h0 = self.shape_code.expand(b, 512, 4, 4, 4).clone()
a0 = self.mlp0(z)
h0 = actvn( adaIN(h0,a0) )
h1 = self.enc_1(h0)
a1 = self.mlp1(z)
h1 = actvn( adaIN(h1,a1) )
h2 = self.enc_2(h1)
a2 = self.mlp2(z)
h2 = actvn(adaIN(h2, a2))
# Rotation
h2_rot = rot(h2,angle,euler_seq=self.euler_seq,padding="border")
b,c,d,h,w = h2_rot.size()
h2_2d = h2_rot.contiguous().view(b,c*d,h,w)
h2_2d = actvn(self.proj(h2_2d))
# Upsampling 2D
h3 = self.enc_3(h2_2d)
a3 = self.mlp3(z)
h3 = actvn(adaIN(h3, a3))
h4 = self.enc_4(h3)
a4 = self.mlp4(z)
h4 = actvn(adaIN(h4, a4))
h5 = self.enc_5(h4)
return F.tanh(h5)
def actvn(x):
out = F.leaky_relu(x, 2e-1)
return out
class Encoder(nn.Module):
def __init__(self,in_dim=3, in_size=64, z_dim=128):
super().__init__()
self.model = nn.Sequential(*[
nn.Conv2d( in_dim, 64, 3, 2, 1), nn.LeakyReLU(0.2),
nn.Conv2d( 64, 128, 3, 2, 1), nn.LeakyReLU(0.2), nn.InstanceNorm2d(128),
nn.Conv2d(128, 256, 3, 2, 1), nn.LeakyReLU(0.2), nn.InstanceNorm2d(256),
nn.Conv2d(256, 512, 3, 2, 1), nn.LeakyReLU(0.2), nn.InstanceNorm2d(512),
])
self.enc_out = nn.Sequential(*[
nn.Linear((in_size//16)**2*512,128), nn.LeakyReLU(0.2),
nn.Linear(128,z_dim), nn.Tanh()
])
def forward(self, x):
b,c,h,w = x.shape
x = self.model.forward(x).view(b,(h//16)**2*512)
enc = self.enc_out(x)
return enc
class LinearBlock(nn.Module):
def __init__(self, input_dim, output_dim, norm='none', activation='relu'):
super(LinearBlock, self).__init__()
use_bias = True
# initialize fully connected layer
self.fc = nn.Linear(input_dim, output_dim, bias=use_bias)
# initialize normalization
norm_dim = output_dim
if norm == 'bn':
self.norm = nn.BatchNorm1d(norm_dim)
elif norm == 'in':
self.norm = nn.InstanceNorm1d(norm_dim)
elif norm == 'none' or norm == 'sn':
self.norm = None
else:
assert 0, "Unsupported normalization: {}".format(norm)
# initialize activation
if activation == 'relu':
self.activation = nn.ReLU(inplace=True)
elif activation == 'lrelu':
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == 'prelu':
self.activation = nn.PReLU()
elif activation == 'selu':
self.activation = nn.SELU(inplace=True)
elif activation == 'tanh':
self.activation = nn.Tanh()
elif activation == 'none':
self.activation = None
else:
assert 0, "Unsupported activation: {}".format(activation)
def forward(self, x):
out = self.fc(x)
if self.norm:
out = self.norm(out)
if self.activation:
out = self.activation(out)
return out
def rot(x,angle,euler_seq='xyz',padding='zeros'):
b,c,d,h,w = x.shape
grid = set_id_grid(x)
grid_flat = grid.reshape(b, 3, -1)
grid_rot_flat = euler2mat(angle,euler_seq=euler_seq).bmm(grid_flat)
grid_rot = grid_rot_flat.reshape(b,3,d,h,w)
x_rot = F.grid_sample(x,grid_rot.permute(0,2,3,4,1),padding_mode=padding,mode='bilinear')
return x_rot
def euler2mat(angle, euler_seq='xyz' ):
"""Convert euler angles to rotation matrix.
Reference: https://github.com/pulkitag/pycaffe-utils/blob/master/rot_utils.py#L174
Args:
angle: rotation angle along 3 axis (in radians) -- size = [B, 3]
Returns:
Rotation matrix corresponding to the euler angles -- size = [B, 3, 3]
"""
B = angle.size(0)
x, y, z = angle[:,0], angle[:,1], angle[:,2]
zeros = z.detach()*0
ones = zeros.detach()+1
cosz = torch.cos(z)
sinz = torch.sin(z)
zmat = torch.stack([cosz, -sinz, zeros,
sinz, cosz, zeros,
zeros, zeros, ones], dim=1).reshape(B, 3, 3)
cosy = torch.cos(y)
siny = torch.sin(y)
ymat = torch.stack([cosy, zeros, siny,
zeros, ones, zeros,
-siny, zeros, cosy], dim=1).reshape(B, 3, 3)
cosx = torch.cos(x)
sinx = torch.sin(x)
xmat = torch.stack([ones, zeros, zeros,
zeros, cosx, -sinx,
zeros, sinx, cosx], dim=1).reshape(B, 3, 3)
if euler_seq == 'xyz':
rotMat = xmat.bmm(ymat).bmm(zmat)
elif euler_seq == 'zyx':
rotMat = zmat.bmm(ymat).bmm(xmat)
return rotMat
def set_id_grid(x):
b, c, d, h, w = x.shape
z_range = (torch.linspace(-1,1,steps=d)).view(1, d, 1, 1).expand(1, d, h, w).type_as(x) # [1, H, W, D]
y_range = (torch.linspace(-1,1,steps=h)).view(1, 1, h, 1).expand(1, d, h, w).type_as(x) # [1, H, W, D]
x_range = (torch.linspace(-1,1,steps=w)).view(1, 1, 1, w).expand(1, d, h, w).type_as(x) # [1, H, W, D]
grid = torch.cat((x_range, y_range, z_range), dim=0)[None,...] # x,y,z
grid = grid.expand(b,3,d,h,w)
return grid
def calc_mean_std(feat, eps=1e-5):
# eps is a small value added to the variance to avoid divide-by-zero.
size = feat.size()
assert (len(size) == 4 or len(size) == 5)
N, C = size[:2]
feat_var = feat.view(N, C, -1).var(dim=2) + eps
feat_std = feat_var.sqrt().view(N, C, 1, 1)
feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
if len(size)==5:
feat_std = feat_std.unsqueeze(-1)
feat_mean = feat_mean.unsqueeze(-1)
return feat_mean, feat_std
def adaIN(content_feat, style_mean_std):
assert(content_feat.size(1) == style_mean_std.size(1)/2)
size = content_feat.size()
b,c = style_mean_std.size()
style_mean, style_std = style_mean_std[:,:c//2],style_mean_std[:,c//2:]
style_mean = style_mean.unsqueeze(-1).unsqueeze(-1)
style_std = style_std.unsqueeze(-1).unsqueeze(-1)
if len(size)==5:
style_mean = style_mean.unsqueeze(-1)
style_std = style_std.unsqueeze(-1)
content_mean, content_std = calc_mean_std(content_feat)
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
def add_SN(m):
for name, c in m.named_children():
m.add_module(name, add_SN(c))
if isinstance(m, (nn.Conv2d, nn.Linear)):
return spectral_norm.spectral_norm(m)#nn.utils.spectral_norm(m)
else:
return m
def init_weights(net, init_type='normal', init_gain=0.02):
"""Initialize network weights.
Parameters:
net (network) -- network to be initialized
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
work better for some applications. Feel free to try yourself.
"""
def init_func(m): # define the initialization function
classname = m.__class__.__name__
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
if init_type == 'normal':
init.normal_(m.weight.data, 0.0, init_gain)
elif init_type == 'xavier':
init.xavier_normal_(m.weight.data, gain=init_gain)
elif init_type == 'kaiming':
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
elif init_type == 'orthogonal':
init.orthogonal_(m.weight.data, gain=init_gain)
else:
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
if hasattr(m, 'bias') and m.bias is not None:
init.constant_(m.bias.data, 0.0)
elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
init.normal_(m.weight.data, 1.0, init_gain)
init.constant_(m.bias.data, 0.0)
print('initialize network with %s' % init_type)
net.apply(init_func) # apply the initialization function <init_func>
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
"""Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
Parameters:
net (network) -- the network to be initialized
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
Return an initialized network.
"""
if len(gpu_ids) > 0:
assert(torch.cuda.is_available())
net.to(gpu_ids[0])
# net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs
if init_type is not None:
init_weights(net, init_type, init_gain=init_gain)
return net
def get_scheduler(optimizer, opt):
"""Return a learning rate scheduler
Parameters:
optimizer -- the optimizer of the network
opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions.
opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
For 'linear', we keep the same learning rate for the first <opt.niter> epochs
and linearly decay the rate to zero over the next <opt.niter_decay> epochs.
For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
See https://pytorch.org/docs/stable/optim.html for more details.
"""
if opt.lr_policy == 'linear':
def lambda_rule(epoch):
lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)
return lr_l
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
elif opt.lr_policy == 'step':
scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
elif opt.lr_policy == 'plateau':
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
elif opt.lr_policy == 'cosine':
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
else:
return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
return scheduler
if __name__ == "__main__":
pass