def create_model(num_blocks=3, block_size=48, sym_recon_grad=False, actnorm=False, split_prior=False, recon_loss_weight=1.0): current_size = (3, 32, 32) alpha = 1e-6 layers = [ Dequantization(UniformDistribution(size=current_size)), Normalization(translation=0, scale=256), Normalization(translation=-alpha, scale=1 / (1 - 2 * alpha)), LogitTransform(), ] for l in range(num_blocks): layers.append(Squeeze()) current_size = (current_size[0] * 4, current_size[1] // 2, current_size[2] // 2) for k in range(block_size): if actnorm: layers.append(ActNorm(current_size[0])) layers.append(Conv1x1(current_size[0])) layers.append(Coupling(current_size)) if split_prior and l < num_blocks - 1: layers.append(SplitPrior(current_size, NegativeGaussianLoss)) current_size = (current_size[0] // 2, current_size[1], current_size[2]) return FlowSequential(NegativeGaussianLoss(size=current_size), *layers)
def create_model(num_blocks=2, block_size=16, sym_recon_grad=False, actnorm=False, split_prior=False, recon_loss_weight=100.0): alpha = 1e-6 layers = [ Dequantization(UniformDistribution(size=(1, 28, 28))), Normalization(translation=0, scale=256), Normalization(translation=-alpha, scale=1 / (1 - 2 * alpha)), LogitTransform(), ] current_size = (1, 28, 28) for l in range(num_blocks): layers.append(Squeeze()) current_size = (current_size[0]*4, current_size[1]//2, current_size[2]//2) for k in range(block_size): if actnorm: layers.append(ActNorm(current_size[0])) layers.append(SelfNormConv(current_size[0], current_size[0], (1, 1), bias=True, stride=1, padding=0, sym_recon_grad=sym_recon_grad, recon_loss_weight=recon_loss_weight)) layers.append(Coupling(current_size)) if split_prior and l < num_blocks - 1: layers.append(SplitPrior(current_size, NegativeGaussianLoss)) current_size = (current_size[0] // 2, current_size[1], current_size[2]) return FlowSequential(NegativeGaussianLoss(size=current_size), *layers)
def test_logdet(input_size=(12, 4, 8, 8)): check_logdet(LearnableLeakyRelu().to('cuda'), input_size) check_logdet(SplineActivation(input_size).to('cuda'), input_size) check_logdet(SmoothLeakyRelu().to('cuda'), input_size) check_logdet(SmoothTanh().to('cuda'), input_size) check_logdet(Identity().to('cuda'), input_size) check_logdet(ActNorm(input_size[1]).to('cuda'), input_size) check_logdet(Coupling(input_size[1:]).to('cuda'), input_size) check_logdet( Normalization(translation=-1e-6, scale=1 / (1 - 2 * 1e-6)).to('cuda'), input_size) check_logdet(Squeeze().to('cuda'), input_size) check_logdet(UnSqueeze().to('cuda'), input_size) test_snf_logdet(input_size) print("All log-det tests passed")
def test_inverses(input_size=(12, 4, 16, 16)): check_inverse(LearnableLeakyRelu().to('cuda'), input_size) check_inverse(SplineActivation(input_size).to('cuda'), input_size) check_inverse(SmoothLeakyRelu().to('cuda'), input_size) check_inverse(SmoothTanh().to('cuda'), input_size) check_inverse(Identity().to('cuda'), input_size) check_inverse(ActNorm(input_size[1]).to('cuda'), input_size) check_inverse(Conv1x1(input_size[1]).to('cuda'), input_size) check_inverse(Conv1x1Householder(input_size[1], 10).to('cuda'), input_size) check_inverse(Coupling(input_size[1:]).to('cuda'), input_size) check_inverse( Normalization(translation=-1e-6, scale=1 / (1 - 2 * 1e-6)).to('cuda'), input_size) check_inverse(Squeeze().to('cuda'), input_size) check_inverse(UnSqueeze().to('cuda'), input_size) test_snf_layer_inverses(input_size) print("All inverse tests passed")