def construct_model(num_epochs,mcr,spectral_loss,save_batch_interval):
"""Construct LBANN model.
"""
import lbann
# Layer graph
input = lbann.Input(target_mode='N/A',name='inp_img')
### Create expected labels for real and fake data (with label flipping = 0.01)
prob_flip=0.01
label_flip_rand = lbann.Uniform(min=0,max=1, neuron_dims='1')
label_flip_prob = lbann.Constant(value=prob_flip, num_neurons='1')
ones = lbann.GreaterEqual(label_flip_rand,label_flip_prob, name='is_real')
zeros = lbann.LogicalNot(ones,name='is_fake')
gen_ones=lbann.Constant(value=1.0,num_neurons='1')## All ones: no flip. Input for training Generator.
#==============================================
### Implement GAN
##Create the noise vector
z = lbann.Reshape(lbann.Gaussian(mean=0.0,stdev=1.0, neuron_dims="64", name='noise_vec'),dims='1 64')
## Creating the GAN object and implementing forward pass for both networks ###
d1_real, d1_fake, d_adv, gen_img, img = ExaGAN.CosmoGAN(mcr)(input,z,mcr)
#==============================================
### Compute quantities for adding to Loss and Metrics
d1_real_bce = lbann.SigmoidBinaryCrossEntropy([d1_real,ones],name='d1_real_bce')
d1_fake_bce = lbann.SigmoidBinaryCrossEntropy([d1_fake,zeros],name='d1_fake_bce')
d_adv_bce = lbann.SigmoidBinaryCrossEntropy([d_adv,gen_ones],name='d_adv_bce')
#img_loss = lbann.MeanSquaredError([gen_img,img])
#l1_loss = lbann.L1Norm(lbann.WeightedSum([gen_img,img], scaling_factors="1 -1"))
#==============================================
### Set up source and destination layers
layers = list(lbann.traverse_layer_graph(input))
weights = set()
src_layers,dst_layers = [],[]
for l in layers:
if(l.weights and "disc1" in l.name and "instance1" in l.name):
src_layers.append(l.name)
#freeze weights in disc2, analogous to discrim.trainable=False in Keras
if(l.weights and "disc2" in l.name):
dst_layers.append(l.name)
for idx in range(len(l.weights)):
l.weights[idx].optimizer = lbann.NoOptimizer()
weights.update(l.weights)
#l2_reg = lbann.L2WeightRegularization(weights=weights, scale=1e-4)
#==============================================
### Define Loss and Metrics
#Define loss (Objective function)
loss_list=[d1_real_bce,d1_fake_bce,d_adv_bce] ## Usual GAN loss function
# loss_list=[d1_real_bce,d1_fake_bce] ## skipping adversarial loss for G for testing spectral loss
if spectral_loss:
dft_gen_img = lbann.DFTAbs(gen_img)
dft_img = lbann.StopGradient(lbann.DFTAbs(img))
spec_loss = lbann.Log(lbann.MeanSquaredError(dft_gen_img, dft_img))
loss_list.append(lbann.LayerTerm(spec_loss, scale=8.0))
loss = lbann.ObjectiveFunction(loss_list)
#Define metrics
metrics = [lbann.Metric(d1_real_bce,name='d_real'),lbann.Metric(d1_fake_bce, name='d_fake'), lbann.Metric(d_adv_bce,name='gen_adv')]
if spectral_loss: metrics.append(lbann.Metric(spec_loss,name='spec_loss'))
#==============================================
### Define callbacks list
callbacks_list=[]
dump_outputs=True
save_model=False
print_model=False
callbacks_list.append(lbann.CallbackPrint())
callbacks_list.append(lbann.CallbackTimer())
callbacks_list.append(lbann.CallbackReplaceWeights(source_layers=list2str(src_layers), destination_layers=list2str(dst_layers),batch_interval=1))
if dump_outputs:
#callbacks_list.append(lbann.CallbackDumpOutputs(layers='inp_img gen_img_instance1_activation', execution_modes='train validation', directory='dump_outs',batch_interval=save_batch_interval,format='npy'))
callbacks_list.append(lbann.CallbackDumpOutputs(layers='gen_img_instance1_activation', execution_modes='train validation', directory='dump_outs',batch_interval=save_batch_interval,format='npy'))
if save_model : callbacks_list.append(lbann.CallbackSaveModel(dir='models'))
if print_model: callbacks_list.append(lbann.CallbackPrintModelDescription())
### Construct model
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
metrics=metrics,
objective_function=loss,
callbacks=callbacks_list)
def construct_model():
"""Construct LBANN model.
ExaGAN model
"""
import lbann
# Layer graph
input = lbann.Input(target_mode='N/A',name='inp_img')
#label flipping
label_flip_rand = lbann.Uniform(min=0,max=1, neuron_dims='1')
label_flip_prob = lbann.Constant(value=0.01, num_neurons='1')
one = lbann.GreaterEqual(label_flip_rand,label_flip_prob, name='is_real')
zero = lbann.LogicalNot(one,name='is_fake')
z = lbann.Reshape(lbann.Gaussian(mean=0.0,stdev=1.0, neuron_dims="64", name='noise_vec'),dims='1 64')
d1_real, d1_fake, d_adv, gen_img = ExaGAN.CosmoGAN()(input,z)
d1_real_bce = lbann.SigmoidBinaryCrossEntropy([d1_real,one],name='d1_real_bce')
d1_fake_bce = lbann.SigmoidBinaryCrossEntropy([d1_fake,zero],name='d1_fake_bce')
d_adv_bce = lbann.SigmoidBinaryCrossEntropy([d_adv,one],name='d_adv_bce')
layers = list(lbann.traverse_layer_graph(input))
# Setup objective function
weights = set()
src_layers = []
dst_layers = []
for l in layers:
if(l.weights and "disc1" in l.name and "instance1" in l.name):
src_layers.append(l.name)
#freeze weights in disc2, analogous to discrim.trainable=False in Keras
if(l.weights and "disc2" in l.name):
dst_layers.append(l.name)
for idx in range(len(l.weights)):
l.weights[idx].optimizer = lbann.NoOptimizer()
weights.update(l.weights)
#l2_reg = lbann.L2WeightRegularization(weights=weights, scale=1e-4)
obj = lbann.ObjectiveFunction([d1_real_bce,d1_fake_bce,d_adv_bce])
# Initialize check metric callback
metrics = [lbann.Metric(d1_real_bce,name='d_real'),
lbann.Metric(d1_fake_bce, name='d_fake'),
lbann.Metric(d_adv_bce,name='gen')]
callbacks = [lbann.CallbackPrint(),
lbann.CallbackTimer(),
#Uncomment to dump output for plotting and further statistical analysis
#lbann.CallbackDumpOutputs(layers='inp_img gen_img_instance1_activation',
# execution_modes='train validation',
# directory='dump_outs',
# batch_interval=100,
# format='npy'),
lbann.CallbackReplaceWeights(source_layers=list2str(src_layers),
destination_layers=list2str(dst_layers),
batch_interval=2)]
# Construct model
num_epochs = 20
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
metrics=metrics,
objective_function=obj,
callbacks=callbacks)