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)