def test_surjection_is_well_behaved(self): batch_size = 10 shape = [8, 4, 4] x = torch.randn(batch_size, *shape) surjections = [ (Augment(StandardNormal([1, 4, 4]), x_size=8, split_dim=1), (9, 4, 4)), (Augment(StandardNormal([4, 4, 4]), x_size=8, split_dim=1), (12, 4, 4)), (Augment(StandardNormal([7, 4, 4]), x_size=8, split_dim=1), (15, 4, 4)), (Augment(StandardNormal([8, 2, 4]), x_size=4, split_dim=2), (8, 6, 4)), (Augment(StandardNormal([8, 4, 2]), x_size=4, split_dim=3), (8, 4, 6)), (Augment(ConditionalMeanNormal(nn.Conv2d(8, 1, kernel_size=1)), x_size=8, split_dim=1), (9, 4, 4)), (Augment(ConditionalMeanNormal(nn.Conv2d(8, 4, kernel_size=1)), x_size=8, split_dim=1), (12, 4, 4)), (Augment(ConditionalMeanNormal(nn.Conv2d(8, 7, kernel_size=1)), x_size=8, split_dim=1), (15, 4, 4)) ] for surjection, new_shape in surjections: with self.subTest(surjection=surjection): self.assert_surjection_is_well_behaved(surjection, x, z_shape=(batch_size, *new_shape), z_dtype=x.dtype)
def __init__(self, num_flows, actnorm, affine, scale_fn_str, hidden_units, activation, range_flow, augment_size, base_dist, cond_size): D = 2 # Number of data dimensions A = D + augment_size # Number of augmented data dimensions P = 2 if affine else 1 # Number of elementwise parameters # initialize context. Only upsample context in ContextInit if latent shape doesn't change during the flow. context_init = MLP(input_size=cond_size, output_size=D, hidden_units=hidden_units, activation=activation) # initialize flow with either augmentation or Abs surjection if augment_size > 0: assert augment_size % 2 == 0 transforms = [Augment(StandardNormal((augment_size, )), x_size=D)] else: transforms = [] transforms = [SimpleAbsSurjection()] if range_flow == 'logit': transforms += [ ScaleBijection(scale=torch.tensor([[1 / 4, 1 / 4]])), Logit() ] elif range_flow == 'softplus': transforms += [SoftplusInverse()] # apply coupling layer flows for _ in range(num_flows): net = nn.Sequential( MLP(A // 2 + D, P * A // 2, hidden_units=hidden_units, activation=activation), ElementwiseParams(P)) if affine: transforms.append( ConditionalAffineCouplingBijection( net, scale_fn=scale_fn(scale_fn_str))) else: transforms.append(ConditionalAdditiveCouplingBijection(net)) if actnorm: transforms.append(ActNormBijection(D)) transforms.append(Reverse(A)) transforms.pop() if base_dist == "uniform": base = StandardUniform((A, )) else: base = StandardNormal((A, )) super(SRFlow, self).__init__(base_dist=base, transforms=transforms, context_init=context_init)
return nn.Sequential( DenseNet(in_channels=channels // 2, out_channels=channels, num_blocks=1, mid_channels=64, depth=8, growth=16, dropout=0.0, gated_conv=True, zero_init=True), ElementwiseParams2d(2)) model = Flow(base_dist=StandardNormal((24, 8, 8)), transforms=[ UniformDequantization(num_bits=8), Augment(StandardUniform((3, 32, 32)), x_size=3), AffineCouplingBijection(net(6)), ActNormBijection2d(6), Conv1x1(6), AffineCouplingBijection(net(6)), ActNormBijection2d(6), Conv1x1(6), AffineCouplingBijection(net(6)), ActNormBijection2d(6), Conv1x1(6), AffineCouplingBijection(net(6)), ActNormBijection2d(6), Conv1x1(6), Squeeze2d(), Slice(StandardNormal((12, 16, 16)), num_keep=12), AffineCouplingBijection(net(12)),
def __init__(self, data_shape, cond_shape, num_bits, num_scales, num_steps, actnorm, conditional_channels, lowres_encoder_channels, lowres_encoder_blocks, lowres_encoder_depth, lowres_upsampler_channels, pooling, compression_ratio, coupling_network, coupling_blocks, coupling_channels, coupling_dropout=0.0, coupling_gated_conv=None, coupling_depth=None, coupling_mixtures=None, dequant="flow", dequant_steps=4, dequant_context=32, dequant_blocks=2, augment_steps=4, augment_context=32, augment_blocks=2, augment_size=None, checkerboard_scales=[], tuple_flip=True): if len(compression_ratio) == 1 and num_scales > 1: compression_ratio = [compression_ratio[0]] * (num_scales - 1) assert all([ compression_ratio[s] >= 0.0 and compression_ratio[s] < 1.0 for s in range(num_scales - 1) ]) # initialize context. Only upsample context in ContextInit if latent shape doesn't change during the flow. context_init = ContextInit(num_bits=num_bits, in_channels=cond_shape[0], out_channels=lowres_encoder_channels, mid_channels=lowres_encoder_channels, num_blocks=lowres_encoder_blocks, depth=lowres_encoder_depth, dropout=coupling_dropout) transforms = [] current_shape = data_shape if dequant == 'uniform': transforms.append(UniformDequantization(num_bits=num_bits)) elif dequant == 'flow': dequantize_flow = DequantizationFlow( data_shape=data_shape, num_bits=num_bits, num_steps=dequant_steps, coupling_network=coupling_network, num_context=dequant_context, num_blocks=dequant_blocks, mid_channels=coupling_channels, depth=coupling_depth, dropout=0.0, gated_conv=False, num_mixtures=coupling_mixtures, checkerboard=True, tuple_flip=tuple_flip) transforms.append( VariationalDequantization(encoder=dequantize_flow, num_bits=num_bits)) # Change range from [0,1]^D to [-0.5, 0.5]^D transforms.append(ScalarAffineBijection(shift=-0.5)) # Initial squeezing if current_shape[1] >= 128 and current_shape[2] >= 128: # H x W -> 64 x 64 transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) if current_shape[1] >= 64 and current_shape[2] >= 64: # H x W -> 32 x 32 transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) if 0 not in checkerboard_scales or (current_shape[1] > 32 and current_shape[2] > 32): # Only go to 16 x 16 if not doing checkerboard splits first transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) # add in augmentation channels if desired if augment_size is not None and augment_size > 0: #transforms.append(Augment(StandardUniform((augment_size, current_shape[1], current_shape[2])), x_size=current_shape[0])) #transforms.append(Augment(StandardNormal((augment_size, current_shape[1], current_shape[2])), x_size=current_shape[0])) augment_flow = AugmentFlow(data_shape=current_shape, augment_size=augment_size, num_steps=augment_steps, coupling_network=coupling_network, mid_channels=coupling_channels, num_context=augment_context, num_mixtures=coupling_mixtures, num_blocks=augment_blocks, dropout=0.0, checkerboard=True, tuple_flip=tuple_flip) transforms.append( Augment(encoder=augment_flow, x_size=current_shape[0])) current_shape = (current_shape[0] + augment_size, current_shape[1], current_shape[2]) for scale in range(num_scales): # First and Third scales use checkerboard split pattern checkerboard = scale in checkerboard_scales context_out_channels = min(current_shape[0], coupling_channels) context_out_shape = (context_out_channels, current_shape[1], current_shape[2] // 2) if checkerboard else (context_out_channels, current_shape[1], current_shape[2]) # reshape the context to the current size for all flow steps at this scale context_upsampler_net = UpsamplerNet( in_channels=lowres_encoder_channels, out_shape=context_out_shape, mid_channels=lowres_upsampler_channels) transforms.append( ContextUpsampler(context_net=context_upsampler_net, direction='forward')) for step in range(num_steps): flip = (step % 2 == 0) if tuple_flip else False if len(conditional_channels) == 0: if actnorm: transforms.append(ActNormBijection2d(current_shape[0])) transforms.append(Conv1x1(current_shape[0])) else: if actnorm: transforms.append( ConditionalActNormBijection2d( cond_shape=current_shape, out_channels=current_shape[0], mid_channels=conditional_channels)) transforms.append( ConditionalConv1x1(cond_shape=current_shape, out_channels=current_shape[0], mid_channels=conditional_channels, slogdet_cpu=True)) if coupling_network in ["conv", "densenet"]: transforms.append( SRCoupling(x_size=context_out_shape, y_size=current_shape, mid_channels=coupling_channels, depth=coupling_depth, num_blocks=coupling_blocks, dropout=coupling_dropout, gated_conv=coupling_gated_conv, coupling_network=coupling_network, checkerboard=checkerboard, flip=flip)) elif coupling_network == "transformer": transforms.append( SRMixtureCoupling(x_size=context_out_shape, y_size=current_shape, mid_channels=coupling_channels, dropout=coupling_dropout, num_blocks=coupling_blocks, num_mixtures=coupling_mixtures, checkerboard=checkerboard, flip=flip)) # Upsample context (for the previous flows, only if moving in the inverse direction) transforms.append( ContextUpsampler(context_net=context_upsampler_net, direction='inverse')) if scale < num_scales - 1: if pooling == 'none' or compression_ratio[scale] == 0.0: # fully bijective flow with multi-scale architecture transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) elif pooling == 'slice': # slice some of the dimensions (channel-wise) out from further flow steps unsliced_channels = int( max( 1, 4 * current_shape[0] * (1.0 - compression_ratio[scale]))) sliced_channels = int(4 * current_shape[0] - unsliced_channels) noise_shape = (sliced_channels, current_shape[1] // 2, current_shape[2] // 2) transforms.append(Squeeze2d()) transforms.append( Slice(StandardNormal(noise_shape), num_keep=unsliced_channels, dim=1)) current_shape = (unsliced_channels, current_shape[1] // 2, current_shape[2] // 2) elif pooling == 'max': # max pooling to compress dimensions spatially, h//2 and w//2 noise_shape = (current_shape[0] * 3, current_shape[1] // 2, current_shape[2] // 2) decoder = StandardHalfNormal(noise_shape) transforms.append( SimpleMaxPoolSurjection2d(decoder=decoder)) current_shape = (current_shape[0], current_shape[1] // 2, current_shape[2] // 2) elif pooling == "mvae": # Compressive flow: reduce the dimensionality of data by 2 (channel-wise) compressed_channels = max( 1, int(current_shape[0] * (1.0 - compression_ratio[scale]))) latent_size = compressed_channels * current_shape[ 1] * current_shape[2] vae_channels = [ current_shape[0] * 2, current_shape[0] * 4, current_shape[0] * 8 ] encoder = ConditionalNormal(ConvEncoderNet( in_channels=current_shape[0], out_channels=latent_size, mid_channels=vae_channels, max_pool=True, batch_norm=True), split_dim=1) decoder = ConditionalNormal(ConvDecoderNet( in_channels=latent_size, out_shape=(current_shape[0] * 2, current_shape[1], current_shape[2]), mid_channels=list(reversed(vae_channels)), batch_norm=True, in_lambda=lambda x: x.view(x.shape[0], x.shape[1], 1, 1 )), split_dim=1) transforms.append(VAE(encoder=encoder, decoder=decoder)) transforms.append( Reshape(input_shape=(latent_size, ), output_shape=(compressed_channels, current_shape[1], current_shape[2]))) # after reducing channels with mvae, squeeze to reshape latent space before another sequence of flows transforms.append(Squeeze2d()) current_shape = ( compressed_channels * 4, # current_shape[0] * 4 current_shape[1] // 2, current_shape[2] // 2) else: raise ValueError( "pooling argument must be either mvae, slice, max, or none" ) else: if actnorm: transforms.append(ActNormBijection2d(current_shape[0])) # for reference save the shape output by the bijective flow self.latent_size = current_shape[0] * current_shape[1] * current_shape[ 2] self.flow_shape = current_shape super(SRPoolFlow, self).__init__(base_dist=ConvNormal2d(current_shape), transforms=transforms, context_init=context_init)
## Specify data ## ################## train_loader, test_loader = get_data(args) ################### ## Specify model ## ################### assert args.augdim % 2 == 0 D = 2 # Number of data dimensions A = 2 + args.augdim # Number of augmented data dimensions P = 2 if args.affine else 1 # Number of elementwise parameters transforms = [Augment(StandardNormal((args.augdim, )), x_size=D)] for _ in range(args.num_flows): net = nn.Sequential( MLP(A // 2, P * A // 2, hidden_units=args.hidden_units, activation=args.activation), ElementwiseParams(P)) if args.affine: transforms.append( AffineCouplingBijection(net, scale_fn=scale_fn(args.scale_fn))) else: transforms.append(AdditiveCouplingBijection(net)) if args.actnorm: transforms.append(ActNormBijection(D)) transforms.append(Reverse(A)) transforms.pop()
def __init__(self, num_flows, actnorm, affine, scale_fn_str, hidden_units, activation, range_flow, augment_size, base_dist): D = 2 # Number of data dimensions if base_dist == "uniform": classifier = MLP(D, D // 2, hidden_units=hidden_units, activation=activation, out_lambda=lambda x: x.view(-1)) transforms = [ ElementAbsSurjection(classifier=classifier), ShiftBijection(shift=torch.tensor([[0.0, 4.0]])), ScaleBijection(scale=torch.tensor([[1 / 4, 1 / 8]])) ] base = StandardUniform((D, )) else: A = D + augment_size # Number of augmented data dimensions P = 2 if affine else 1 # Number of elementwise parameters # initialize flow with either augmentation or Abs surjection if augment_size > 0: assert augment_size % 2 == 0 transforms = [ Augment(StandardNormal((augment_size, )), x_size=D) ] else: transforms = [SimpleAbsSurjection()] if range_flow == 'logit': transforms += [ ScaleBijection(scale=torch.tensor([[1 / 4, 1 / 4]])), Logit() ] elif range_flow == 'softplus': transforms += [SoftplusInverse()] # apply coupling layer flows for _ in range(num_flows): net = nn.Sequential( MLP(A // 2, P * A // 2, hidden_units=hidden_units, activation=activation), ElementwiseParams(P)) if affine: transforms.append( AffineCouplingBijection( net, scale_fn=scale_fn(scale_fn_str))) else: transforms.append(AdditiveCouplingBijection(net)) if actnorm: transforms.append(ActNormBijection(D)) transforms.append(Reverse(A)) transforms.pop() base = StandardNormal((A, )) super(UnconditionalFlow, self).__init__(base_dist=base, transforms=transforms)
def reduction_layer(channels, items): return [ *perm_norm_bi(channels), *perm_norm_bi(channels), *perm_norm_bi(channels), Squeeze2d(4), Slice(StandardNormal((channels * 2, items)), num_keep=channels * 2), ] model = Flow( base_dist=StandardNormal((base_channels * (2**5), n_items // (4**4))), transforms=[ UniformDequantization(num_bits=8), Augment(StandardUniform((base_channels * 1, n_items)), x_size=base_channels), *reduction_layer(base_channels * (2**1), n_items // (4**1)), *reduction_layer(base_channels * (2**2), n_items // (4**2)), *reduction_layer(base_channels * (2**3), n_items // (4**3)), *reduction_layer(base_channels * (2**4), n_items // (4**4)), # *reduction_layer(base_channels*(2**5), n_items//(4**4)), *perm_norm_bi(base_channels * (2**5)) # AffineCouplingBijection(net(base_channels*2)), ActNormBijection2d(base_channels*2), Conv1x1(base_channels*2), # AffineCouplingBijection(net(base_channels*2)), ActNormBijection2d(base_channels*2), Conv1x1(base_channels*2), # AffineCouplingBijection(net(base_channels*2)), ActNormBijection2d(base_channels*2), Conv1x1(base_channels*2), # Squeeze2d(), Slice(StandardNormal((base_channels*2, n_items//4)), num_keep=base_channels*2), # AffineCouplingBijection(net(base_channels*2)), ActNormBijection2d(base_channels*2), Conv1x1(base_channels*2), # AffineCouplingBijection(net(base_channels*2)), ActNormBijection2d(base_channels*2), Conv1x1(base_channels*2), # AffineCouplingBijection(net(base_channels*2)), ActNormBijection2d(base_channels*2), Conv1x1(base_channels*2), ]).to(device)
def __init__(self, data_shape, num_bits, num_scales, num_steps, actnorm, pooling, compression_ratio, coupling_network, coupling_blocks, coupling_channels, coupling_dropout=0.0, coupling_gated_conv=None, coupling_depth=None, coupling_mixtures=None, dequant="flow", dequant_steps=4, dequant_context=32, dequant_blocks=2, augment_steps=4, augment_context=32, augment_blocks=2, augment_size=None, checkerboard_scales=[], tuple_flip=True): if len(compression_ratio) == 1 and num_scales > 1: compression_ratio = [compression_ratio[0]] * (num_scales - 1) assert all([ compression_ratio[s] >= 0 and compression_ratio[s] < 1 for s in range(num_scales - 1) ]) transforms = [] current_shape = data_shape if dequant == 'uniform': transforms.append(UniformDequantization(num_bits=num_bits)) elif dequant == 'flow': dequantize_flow = DequantizationFlow( data_shape=data_shape, num_bits=num_bits, num_steps=dequant_steps, coupling_network=coupling_network, num_context=dequant_context, num_blocks=dequant_blocks, mid_channels=coupling_channels, depth=coupling_depth, dropout=0.0, gated_conv=False, num_mixtures=coupling_mixtures, checkerboard=True, tuple_flip=tuple_flip) transforms.append( VariationalDequantization(encoder=dequantize_flow, num_bits=num_bits)) # Change range from [0,1]^D to [-0.5, 0.5]^D transforms.append(ScalarAffineBijection(shift=-0.5)) # Initial squeezing if current_shape[1] >= 128 and current_shape[2] >= 128: # H x W -> 64 x 64 transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) if current_shape[1] >= 64 and current_shape[2] >= 64: # H x W -> 32 x 32 transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) if 0 not in checkerboard_scales or (current_shape[1] > 32 and current_shape[2] > 32): # Only go to 16 x 16 if not doing checkerboard splits first transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) # add in augmentation channels if desired if augment_size is not None and augment_size > 0: augment_flow = AugmentFlow(data_shape=current_shape, augment_size=augment_size, num_steps=augment_steps, coupling_network=coupling_network, mid_channels=coupling_channels, num_context=augment_context, num_mixtures=coupling_mixtures, num_blocks=augment_blocks, dropout=0.0, checkerboard=True, tuple_flip=tuple_flip) transforms.append( Augment(encoder=augment_flow, x_size=current_shape[0])) current_shape = (current_shape[0] + augment_size, current_shape[1], current_shape[2]) for scale in range(num_scales): checkerboard = scale in checkerboard_scales for step in range(num_steps): flip = (step % 2 == 0) if tuple_flip else False if actnorm: transforms.append(ActNormBijection2d(current_shape[0])) transforms.append(Conv1x1(current_shape[0])) if coupling_network == "conv": transforms.append( Coupling(in_channels=current_shape[0], num_blocks=coupling_blocks, mid_channels=coupling_channels, depth=coupling_depth, dropout=coupling_dropout, gated_conv=coupling_gated_conv, coupling_network=coupling_network, checkerboard=checkerboard, flip=flip)) else: transforms.append( MixtureCoupling(in_channels=current_shape[0], mid_channels=coupling_channels, num_mixtures=coupling_mixtures, num_blocks=coupling_blocks, dropout=coupling_dropout, checkerboard=checkerboard, flip=flip)) if scale < num_scales - 1: if pooling in ['bijective', 'none' ] or compression_ratio[scale] == 0.0: transforms.append(Squeeze2d()) current_shape = (current_shape[0] * 4, current_shape[1] // 2, current_shape[2] // 2) elif pooling == 'slice': # slice some of the dimensions (channel-wise) out from further flow steps unsliced_channels = int( max(1, 4 * current_shape[0] * (1.0 - sliced_ratio[scale]))) sliced_channels = int(4 * current_shape[0] - unsliced_channels) noise_shape = (sliced_channels, current_shape[1] // 2, current_shape[2] // 2) transforms.append(Squeeze2d()) transforms.append( Slice(StandardNormal(noise_shape), num_keep=unsliced_channels, dim=1)) current_shape = (unsliced_channels, current_shape[1] // 2, current_shape[2] // 2) elif pooling == 'max': noise_shape = (current_shape[0] * 3, current_shape[1] // 2, current_shape[2] // 2) decoder = StandardHalfNormal(noise_shape) transforms.append( SimpleMaxPoolSurjection2d(decoder=decoder)) current_shape = (current_shape[0], current_shape[1] // 2, current_shape[2] // 2) else: raise ValueError( f"Pooling argument must be either slice, max or none, not: {pooling}" ) else: if actnorm: transforms.append(ActNormBijection2d(current_shape[0])) # for reference save the shape output by the bijective flow self.flow_shape = current_shape self.latent_size = current_shape[0] * current_shape[1] * current_shape[ 2] super(PoolFlow, self).__init__(base_dist=ConvNormal2d(current_shape), transforms=transforms)