def construct_model():
"""Construct LBANN model.
Pilot1 Combo model
"""
import lbann
# Layer graph
data = lbann.Input(data_field='samples')
responses = lbann.Input(data_field='responses')
pred = combo.Combo()(data)
mse = lbann.MeanSquaredError([responses, pred])
SS_res = lbann.Reduction(lbann.Square(lbann.Subtract(responses, pred)),
mode='sum')
#SS_tot = var(x) = mean((x-mean(x))^2)
mini_batch_size = lbann.MiniBatchSize()
mean = lbann.Divide(lbann.BatchwiseReduceSum(responses), mini_batch_size)
SS_tot = lbann.Divide(
lbann.BatchwiseReduceSum(lbann.Square(lbann.Subtract(responses,
mean))),
mini_batch_size)
eps = lbann.Constant(value=1e-07, hint_layer=SS_tot)
r2 = lbann.Subtract(lbann.Constant(value=1, num_neurons='1'),
lbann.Divide(SS_res, lbann.Add(SS_tot, eps)))
metrics = [lbann.Metric(mse, name='mse')]
metrics.append(lbann.Metric(r2, name='r2'))
callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()]
# Construct model
num_epochs = 100
layers = list(lbann.traverse_layer_graph([data, responses]))
return lbann.Model(num_epochs,
layers=layers,
metrics=metrics,
objective_function=mse,
callbacks=callbacks)
def make_model(
data_dim,
latent_dim,
num_epochs,
):
# Layer graph
data = lbann.Input(data_field='samples')
autoencoder = model.autoencoder.FullyConnectedAutoencoder(
data_dim,
latent_dim,
)
reconstructed = autoencoder(data)
loss = lbann.MeanSquaredError(data, reconstructed)
# Metrics
metrics = [
lbann.Metric(loss, name='loss'),
]
# Callbacks
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackDumpWeights(directory='weights',
epoch_interval=num_epochs),
]
# Contruct model
return lbann.Model(
num_epochs,
layers=lbann.traverse_layer_graph(loss),
objective_function=loss,
metrics=metrics,
callbacks=callbacks,
)
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)
# Default data reader
cur_dir = dirname(abspath(__file__))
data_reader_prototext = join(cur_dir, 'data', args.reader_prototext)
print("DATA READER ", data_reader_prototext)
images = lbann.Input(data_field='samples')
responses = lbann.Input(data_field='responses')
num_labels = 5
images = lbann.Reshape(images, dims='1 300 300')
pred = model.PROBIESNet(num_labels)(images)
mse = lbann.MeanSquaredError([responses, pred])
# Pearson Correlation
# rho(x,y) = covariance(x,y) / sqrt( variance(x) * variance(y) )
pearson_r_cov = lbann.Covariance([pred, responses], name="pearson_r_cov")
pearson_r_var1 = lbann.Variance(responses, name="pearson_r_var1")
pearson_r_var2 = lbann.Variance(pred, name="pearson_r_var2")
pearson_r_mult = lbann.Multiply([pearson_r_var1, pearson_r_var2],
name="pearson_r_mult")
pearson_r_sqrt = lbann.Sqrt(pearson_r_mult, name="pearson_r_sqrt")
eps = lbann.Constant(value=1e-07, hint_layer=pearson_r_sqrt)
parser,
default_optimizer="adam",
default_learning_rate=0.001,
)
args = parser.parse_args()
# Construct layer graph
universes = lbann.Input(data_field='samples')
secrets = lbann.Input(data_field='responses')
statistics_group_size = 1 if args.local_batchnorm else -1
preds = CosmoFlow(
input_width=args.input_width,
output_size=args.num_secrets,
use_bn=args.use_batchnorm,
bn_statistics_group_size=statistics_group_size)(universes)
mse = lbann.MeanSquaredError([preds, secrets])
obj = lbann.ObjectiveFunction([mse])
layers = list(lbann.traverse_layer_graph([universes, secrets]))
# Set parallel_strategy
parallel_strategy = get_parallel_strategy_args(
sample_groups=args.mini_batch_size, depth_groups=args.depth_groups)
pooling_id = 0
dropout_id = 0
for i, layer in enumerate(layers):
if layer == secrets:
continue
layer_name = layer.__class__.__name__
if layer_name == 'Pooling':
pooling_id += 1
decode1 = lbann.FullyConnected(decode2neuron,
name="decode1",
has_bias=True,
hint_layer=image)
# Reconstruction error
reconstruction = lbann.Sigmoid(decode1, name="reconstruction")
bin_cross_entropy = lbann.SigmoidBinaryCrossEntropy([decode1, image],
name="bin_cross_entropy")
bin_cross_entropy_sum = lbann.Reduction(bin_cross_entropy,
name="bin_cross_entropy_sum",
mode="sum")
mean_squared_error = lbann.MeanSquaredError([reconstruction, image],
name="mean_squared_error")
layer_list = list(lbann.traverse_layer_graph(input_))
# Set up objective function
layer_term1 = lbann.LayerTerm(bin_cross_entropy)
layer_term2 = lbann.LayerTerm(kldiv)
l2_reg = lbann.L2WeightRegularization(scale=0.0005)
obj = lbann.ObjectiveFunction([layer_term1, layer_term2, l2_reg])
# Metrics
metrics = [lbann.Metric(mean_squared_error, name="mean squared error")]
# Callbacks
callbacks = [
lbann.CallbackPrint(),
args = parser.parse_args()
# ----------------------------------
# Construct layer graph
# ----------------------------------
# Input data
input = lbann.Input(io_buffer='partitioned', target_mode='regression')
universes = lbann.Identity(input)
secrets = lbann.Identity(input)
# CosmoFlow
x = CosmoFlow(args.output_size, args.input_width).forward(universes)
# Loss function
loss = lbann.MeanSquaredError([x, secrets])
# Metrics
metrics = [lbann.Metric(loss, name="MSE", unit="")]
# Callbacks
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackPolyLearningRate(
power=1.0,
num_epochs=100, # TODO: Warn if args.epochs < 100
),
lbann.CallbackGPUMemoryUsage(),
lbann.CallbackDumpOutputs(directory="dump_acts/",
layers="cosmoflow_module1_fc3_instance1 layer3",
def gen_layers(latent_dim, number_of_atoms):
''' Generates the model for the 3D Convolutional Auto Encoder.
returns the Directed Acyclic Graph (DAG) that the lbann
model will run on.
'''
input_ = lbann.Input(target_mode="reconstruction")
tensors = lbann.Identity(input_)
tensors = lbann.Reshape(tensors, dims="11 32 32 32", name="Sample")
# Input tensor shape is (number_of_atoms)x32x32x32
# Encoder
x = lbann.Identity(tensors)
for i in range(4):
out_channels = latent_dim // (2**(3 - i))
x = lbann.Convolution(x,
num_dims=3,
num_output_channels=out_channels,
num_groups=1,
conv_dims_i=4,
conv_strides_i=2,
conv_dilations_i=1,
conv_pads_i=1,
has_bias=True,
name="Conv_{0}".format(i))
x = lbann.BatchNormalization(x, name="Batch_NORM_{0}".format(i + 1))
x = lbann.LeakyRelu(x, name="Conv_{0}_Activation".format(i + 1))
# Shape: (latent_dim)x2x2x2
encoded = lbann.Convolution(x,
num_dims=3,
num_output_channels=latent_dim,
num_groups=1,
conv_dims_i=2,
conv_strides_i=2,
conv_dilations_i=1,
conv_pads_i=0,
has_bias=True,
name="encoded")
# Shape: (latent_dim)1x1x1
# Decoder
x = lbann.Deconvolution(encoded,
num_dims=3,
num_output_channels=number_of_atoms * 16,
num_groups=1,
conv_dims_i=4,
conv_pads_i=0,
conv_strides_i=2,
conv_dilations_i=1,
has_bias=True,
name="Deconv_1")
x = lbann.BatchNormalization(x, name="BN_D1")
x = lbann.Tanh(x, name="Deconv_1_Activation")
for i in range(3):
out_channels = number_of_atoms * (2**(2 - i))
x = lbann.Deconvolution(x,
num_dims=3,
num_output_channels=out_channels,
num_groups=1,
conv_dims_i=4,
conv_pads_i=1,
conv_strides_i=2,
conv_dilations_i=1,
has_bias=True,
name="Deconv_{0}".format(i + 2))
x = lbann.BatchNormalization(x, name="BN_D{0}".format(i + 2))
if (
i != 2
): #Save the last activation layer because we want to dump the outputs
x = lbann.Tanh(x, name="Deconv_{0}_Activation".format(i + 2))
decoded = lbann.Tanh(x, name="decoded")
img_loss = lbann.MeanSquaredError([decoded, tensors])
metrics = [lbann.Metric(img_loss, name='recon_error')]
# ----------------------------------
# Set up DAG
# ----------------------------------
layers = lbann.traverse_layer_graph(input_) #Generate Model DAG
return layers, img_loss, metrics
relu9,
name="decode1",
data_layout="model_parallel",
hint_layer=data,
#num_neurons_from_data_reader=True,
has_bias=True)
relu10 = lbann.Relu(decode1, name="relu10", data_layout="model_parallel")
# Reconstruction layers
reconstruction = lbann.Split(relu10,
name="reconstruction",
data_layout="model_parallel")
mean_squared_error = lbann.MeanSquaredError([reconstruction, data],
name="mean_squared_error",
data_layout="model_parallel")
# Pearson Correlation
# rho(x,y) = covariance(x,y) / sqrt( variance(x) * variance(y) )
pearson_r_cov = lbann.Covariance([reconstruction, data],
name="pearson_r_cov",
data_layout="model_parallel")
pearson_r_var1 = lbann.Variance(data,
name="pearson_r_var1",
data_layout="model_parallel")
pearson_r_var2 = lbann.Variance(reconstruction,
name="pearson_r_var1",
def construct_macc_surrogate_model(xdim, ydim, zdim, wae_mcf, surrogate_mcf,
lambda_cyc, useCNN, dump_models,
pretrained_dir, ltfb_batch_interval,
num_epochs):
"""Construct MACC surrogate model.
See https://arxiv.org/pdf/1912.08113.pdf model architecture and other details
"""
# Layer graph
input = lbann.Input(data_field='samples', name='inp_data')
# data is 64*64*4 images + 15 scalar + 5 param
inp_slice = lbann.Slice(input,
axis=0,
slice_points=str_list([0, ydim, ydim + xdim]),
name='inp_slice')
gt_y = lbann.Identity(inp_slice, name='gt_y')
gt_x = lbann.Identity(inp_slice, name='gt_x') #param not used
zero = lbann.Constant(value=0.0, num_neurons='1', name='zero')
one = lbann.Constant(value=1.0, num_neurons='1', name='one')
z = lbann.Gaussian(mean=0.0, stdev=1.0, neuron_dims="20")
wae = macc_network_architectures.MACCWAE(
zdim, ydim, cf=wae_mcf, use_CNN=useCNN) #pretrained, freeze
inv = macc_network_architectures.MACCInverse(xdim, cf=surrogate_mcf)
fwd = macc_network_architectures.MACCForward(zdim, cf=surrogate_mcf)
y_pred_fwd = wae.encoder(gt_y)
param_pred_ = wae.encoder(gt_y)
input_fake = inv(param_pred_)
output_cyc = fwd(input_fake)
y_image_re2 = wae.decoder(output_cyc)
'''**** Train cycleGAN input params <--> latent space of (images, scalars) ****'''
output_fake = fwd(gt_x)
y_image_re = wae.decoder(output_fake)
param_pred2_ = wae.encoder(y_image_re)
input_cyc = inv(param_pred2_)
L_l2_x = lbann.MeanSquaredError(input_fake, gt_x)
L_cyc_x = lbann.MeanSquaredError(input_cyc, gt_x)
L_l2_y = lbann.MeanSquaredError(output_fake, y_pred_fwd)
L_cyc_y = lbann.MeanSquaredError(output_cyc, y_pred_fwd)
#@todo slice here to separate scalar from image
img_sca_loss = lbann.MeanSquaredError(y_image_re, gt_y)
#L_cyc = L_cyc_y + L_cyc_x
L_cyc = lbann.Add(L_cyc_y, L_cyc_x)
#loss_gen0 = L_l2_y + lamda_cyc*L_cyc
loss_gen0 = lbann.WeightedSum([L_l2_y, L_cyc],
scaling_factors=f'1 {lambda_cyc}')
loss_gen1 = lbann.WeightedSum([L_l2_x, L_cyc_y],
scaling_factors=f'1 {lambda_cyc}')
#loss_gen1 = L_l2_x + lamda_cyc*L_cyc_y
layers = list(lbann.traverse_layer_graph(input))
weights = set()
#Freeze appropriate (pretrained) weights
pretrained_models = ["wae"] #add macc?
for l in layers:
for idx in range(len(pretrained_models)):
if (l.weights and pretrained_models[idx] in l.name):
for w in range(len(l.weights)):
l.weights[w].optimizer = lbann.NoOptimizer()
weights.update(l.weights)
l2_reg = lbann.L2WeightRegularization(weights=weights, scale=1e-4)
#d_adv_bce = lbann.LayerTerm(d_adv_bce,scale=0.01)
# Setup objective function
obj = lbann.ObjectiveFunction([loss_gen0, loss_gen1, l2_reg])
# Initialize check metric callback
metrics = [
lbann.Metric(img_sca_loss, name='fw_loss'),
lbann.Metric(L_l2_x, name='inverse loss'),
lbann.Metric(L_cyc_y, name='output cycle loss'),
lbann.Metric(L_cyc_x, name='param cycle loss')
]
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackSaveModel(dir=dump_models),
lbann.CallbackLoadModel(dirs=str(pretrained_dir)),
lbann.CallbackTimer()
]
if (ltfb_batch_interval > 0):
callbacks.append(
lbann.CallbackLTFB(batch_interval=ltfb_batch_interval,
metric='fw_loss',
low_score_wins=True,
exchange_hyperparameters=True))
# Construct model
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
metrics=metrics,
objective_function=obj,
callbacks=callbacks)
def construct_jag_wae_model(ydim, zdim, mcf, useCNN, dump_models,
ltfb_batch_interval, num_epochs):
"""Construct LBANN model.
JAG Wasserstein autoencoder model
"""
# Layer graph
input = lbann.Input(data_field='samples', name='inp_data')
# data is 64*64*4 images + 15 scalar + 5 param
#inp_slice = lbann.Slice(input, axis=0, slice_points="0 16399 16404",name='inp_slice')
inp_slice = lbann.Slice(input,
axis=0,
slice_points=str_list([0, ydim, ydim + 5]),
name='inp_slice')
gt_y = lbann.Identity(inp_slice, name='gt_y')
gt_x = lbann.Identity(inp_slice, name='gt_x') #param not used
zero = lbann.Constant(value=0.0, num_neurons='1', name='zero')
one = lbann.Constant(value=1.0, num_neurons='1', name='one')
z_dim = 20 #Latent space dim
z = lbann.Gaussian(mean=0.0, stdev=1.0, neuron_dims="20")
model = macc_network_architectures.MACCWAE(zdim,
ydim,
cf=mcf,
use_CNN=useCNN)
d1_real, d1_fake, d_adv, pred_y = model(z, gt_y)
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')
img_loss = lbann.MeanSquaredError([pred_y, gt_y])
rec_error = lbann.L2Norm2(
lbann.WeightedSum([pred_y, gt_y], scaling_factors="1 -1"))
layers = list(lbann.traverse_layer_graph(input))
# Setup objective function
weights = set()
src_layers = []
dst_layers = []
for l in layers:
if (l.weights and "disc0" in l.name and "instance1" in l.name):
src_layers.append(l.name)
#freeze weights in disc2
if (l.weights and "disc1" 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)
d_adv_bce = lbann.LayerTerm(d_adv_bce, scale=0.01)
obj = lbann.ObjectiveFunction(
[d1_real_bce, d1_fake_bce, d_adv_bce, img_loss, rec_error, l2_reg])
# Initialize check metric callback
metrics = [lbann.Metric(img_loss, name='recon_error')]
#pred_y = macc_models.MACCWAE.pred_y_name
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackPrintModelDescription(),
lbann.CallbackSaveModel(dir=dump_models),
lbann.CallbackReplaceWeights(source_layers=list2str(src_layers),
destination_layers=list2str(dst_layers),
batch_interval=2)
]
if (ltfb_batch_interval > 0):
callbacks.append(
lbann.CallbackLTFB(batch_interval=ltfb_batch_interval,
metric='recon_error',
low_score_wins=True,
exchange_hyperparameters=True))
# Construct model
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
metrics=metrics,
objective_function=obj,
callbacks=callbacks)
name="decode1",
data_layout="model_parallel",
hint_layer=image,
has_bias=True)
reconstruction = lbann.Sigmoid(decode1,
name="reconstruction",
data_layout="model_parallel")
dropout2 = lbann.Dropout(reconstruction,
name="dropout2",
data_layout="model_parallel",
keep_prob=0.8)
# Reconstruction
mean_squared_error = lbann.MeanSquaredError([dropout2, image],
name="mean_squared_error")
layer_term = lbann.LayerTerm(mean_squared_error)
obj = lbann.ObjectiveFunction(layer_term)
metrics = [lbann.Metric(mean_squared_error, name=mean_squared_error.name)]
img_strategy = lbann.TrackSampleIDsStrategy(input_layer_name=input_.name,
num_tracked_images=10)
summarize_images = lbann.CallbackSummarizeImages(
selection_strategy=img_strategy,
image_source_layer_name=reconstruction.name,
epoch_interval=1)
# Dump original image from input layer one time
def construct_model():
"""Construct LBANN model.
JAG Wasserstein autoencoder model
"""
import lbann
# Layer graph
input = lbann.Input(target_mode='N/A',name='inp_data')
# data is 64*64*4 images + 15 scalar + 5 param
inp_slice = lbann.Slice(input, axis=0, slice_points="0 16399 16404",name='inp_slice')
gt_y = lbann.Identity(inp_slice,name='gt_y')
gt_x = lbann.Identity(inp_slice, name='gt_x') #param not used
zero = lbann.Constant(value=0.0,num_neurons='1',name='zero')
one = lbann.Constant(value=1.0,num_neurons='1',name='one')
y_dim = 16399 #image+scalar shape
z_dim = 20 #Latent space dim
z = lbann.Gaussian(mean=0.0,stdev=1.0, neuron_dims="20")
d1_real, d1_fake, d_adv, pred_y = jag_models.WAE(z_dim,y_dim)(z,gt_y)
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')
img_loss = lbann.MeanSquaredError([pred_y,gt_y])
rec_error = lbann.L2Norm2(lbann.WeightedSum([pred_y,gt_y], scaling_factors="1 -1"))
layers = list(lbann.traverse_layer_graph(input))
# Setup objective function
weights = set()
src_layers = []
dst_layers = []
for l in layers:
if(l.weights and "disc0" in l.name and "instance1" in l.name):
src_layers.append(l.name)
#freeze weights in disc2
if(l.weights and "disc1" 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)
d_adv_bce = lbann.LayerTerm(d_adv_bce,scale=0.01)
obj = lbann.ObjectiveFunction([d1_real_bce,d1_fake_bce,d_adv_bce,img_loss,rec_error,l2_reg])
# Initialize check metric callback
metrics = [lbann.Metric(img_loss, name='recon_error')]
callbacks = [lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackReplaceWeights(source_layers=list2str(src_layers),
destination_layers=list2str(dst_layers),
batch_interval=2)]
# Construct model
num_epochs = 100
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
metrics=metrics,
objective_function=obj,
callbacks=callbacks)
def construct_model(args):
"""Construct LBANN for CosmoGAN 3D model.
"""
obj = []
metrics = []
callbacks = []
w = [args.input_width]*3
w.insert(0,args.input_channel)
_sample_dims = w
ps = None
#have model and input ps
if(args.use_distconv):
ps = get_parallel_strategy_args(
sample_groups=args.mini_batch_size,
depth_groups=args.depth_groups,
height_groups=args.height_groups,
)
g_device = 'GPU'
input_ = lbann.Input(name='input', data_field='samples')
input_ = lbann.Reshape(input_, dims=list2str(_sample_dims),name='in_reshape', device=g_device),
x1 = lbann.Identity(input_, parallel_strategy=None, name='x1')
x2 = lbann.Identity(input_, name='x2') if args.compute_mse else None
zero = lbann.Constant(value=0.0,num_neurons='1',name='zero',device=g_device)
one = lbann.Constant(value=1.0,num_neurons='1',name='one', device=g_device)
z = lbann.Reshape(lbann.Gaussian(mean=0.0,stdev=1.0, neuron_dims="64", name='noise_vec', device=g_device),
dims='1 64', name='noise_vec_reshape',device=g_device)
print("RUN ARGS ", args)
d1_real,d1_fake,d_adv, gen_img = model.Exa3DGAN(args.input_width,args.input_channel,
g_device,ps,use_bn=args.use_bn)(x1,z)
layers=list(lbann.traverse_layer_graph([d1_real, d1_fake]))
# 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)
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')
mse = lbann.MeanSquaredError([gen_img, x2], name='MSE') if args.compute_mse else None
obj.append(d1_real_bce)
obj.append(d1_fake_bce)
obj.append(d_adv_bce)
metrics.append(lbann.Metric(d_adv_bce, name='d_adv_bce'))
metrics.append(lbann.Metric(d1_real_bce, name='d1_real_bce'))
metrics.append(lbann.Metric(d1_fake_bce, name='d1_fake_bce'))
if (mse is not None):
obj.append(mse)
metrics.append(lbann.Metric(mse, name='MSE'))
callbacks.append(lbann.CallbackPrint())
callbacks.append(lbann.CallbackTimer())
callbacks.append(lbann.CallbackGPUMemoryUsage())
# ------------------------------------------
# Construct model
# ------------------------------------------
return lbann.Model(args.num_epochs,
weights=weights,
layers=layers,
objective_function=obj,
metrics=metrics,
callbacks=callbacks)
def construct_model():
"""Construct MACC surrogate model.
See https://arxiv.org/pdf/1912.08113.pdf model architecture and other details
"""
import lbann
# Layer graph
input = lbann.Input(data_field='samples', name='inp_data')
# data is 64*64*4 images + 15 scalar + 5 param
inp_slice = lbann.Slice(input,
axis=0,
slice_points=str_list(
[0, args.ydim, args.ydim + args.xdim]),
name='inp_slice')
gt_y = lbann.Identity(inp_slice, name='gt_y')
gt_x = lbann.Identity(inp_slice, name='gt_x') #param not used
zero = lbann.Constant(value=0.0, num_neurons='1', name='zero')
one = lbann.Constant(value=1.0, num_neurons='1', name='one')
z = lbann.Gaussian(mean=0.0, stdev=1.0, neuron_dims="20")
wae = macc_models.MACCWAE(args.zdim,
args.ydim,
cf=args.wae_mcf,
use_CNN=args.useCNN) #pretrained, freeze
inv = macc_models.MACCInverse(args.xdim, cf=args.surrogate_mcf)
fwd = macc_models.MACCForward(args.zdim, cf=args.surrogate_mcf)
y_pred_fwd = wae.encoder(gt_y)
param_pred_ = wae.encoder(gt_y)
input_fake = inv(param_pred_)
output_cyc = fwd(input_fake)
y_image_re2 = wae.decoder(output_cyc)
'''**** Train cycleGAN input params <--> latent space of (images, scalars) ****'''
output_fake = fwd(gt_x)
y_image_re = wae.decoder(output_fake)
y_out = wae.decoder(y_pred_fwd)
param_pred2_ = wae.encoder(y_image_re)
input_cyc = inv(param_pred2_)
L_l2_x = lbann.MeanSquaredError(
input_fake, gt_x) #(x,inv(enc(y)), (encoder+)inverse loss
L_cyc_x = lbann.MeanSquaredError(
input_cyc, gt_x) #param, x cycle loss, from latent space
L_l2_y = lbann.MeanSquaredError(
output_fake, y_pred_fwd) #pred error into latent space (enc(y),fw(x))
L_cyc_y = lbann.MeanSquaredError(
output_cyc,
y_pred_fwd) # pred error into latent space (enc(y), fw(inv(enc(y))))
#@todo slice here to separate scalar from image
img_sca_loss = lbann.MeanSquaredError(
y_image_re,
gt_y) # (y,dec(fw(x))) #forward model to decoder, no latent space
dec_fw_inv_enc_y = lbann.MeanSquaredError(
y_image_re2, gt_y) #(y, dec(fw(inv(enc(y))))) y->enc_z->x'->fw_z->y'
wae_loss = lbann.MeanSquaredError(y_out, gt_y) #(y, dec(enc(y)) '
#L_cyc = L_cyc_y + L_cyc_x
L_cyc = lbann.Add(L_cyc_y, L_cyc_x)
#loss_gen0 = L_l2_y + lamda_cyc*L_cyc
loss_gen0 = lbann.WeightedSum([L_l2_y, L_cyc],
scaling_factors=f'1 {args.lamda_cyc}')
loss_gen1 = lbann.WeightedSum([L_l2_x, L_cyc_y],
scaling_factors=f'1 {args.lamda_cyc}')
#loss_gen1 = L_l2_x + lamda_cyc*L_cyc_y
conc_out = lbann.Concatenation(
[gt_x, wae_loss, img_sca_loss, dec_fw_inv_enc_y, L_l2_x],
name='x_errors')
layers = list(lbann.traverse_layer_graph(input))
weights = set()
for l in layers:
weights.update(l.weights)
# Setup objective function
obj = lbann.ObjectiveFunction([loss_gen0, loss_gen1])
# Initialize check metric callback
metrics = [
lbann.Metric(img_sca_loss, name='img_re1'),
lbann.Metric(dec_fw_inv_enc_y, name='img_re2'),
lbann.Metric(wae_loss, name='wae_loss'),
lbann.Metric(L_l2_x, name='inverse loss'),
lbann.Metric(L_cyc_y, name='output cycle loss'),
lbann.Metric(L_cyc_x, name='param cycle loss')
]
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackDumpOutputs(layers=f'{conc_out.name}',
execution_modes='test',
directory=args.dump_outputs,
batch_interval=1,
format='npy'),
lbann.CallbackTimer()
]
# Construct model
num_epochs = 1
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
serialize_io=True,
metrics=metrics,
objective_function=obj,
callbacks=callbacks)
