parameters: ()
return_statement: return lbann.Model(num_epochs,
                         weights=weights,
                         layers=layers,
                         serialize_io=True,
                         metrics=metrics,
                         objective_function=obj,
                         callbacks=callbacks)
docstring: Construct MACC surrogate model.
  
      See https://arxiv.org/pdf/1912.08113.pdf model architecture and other details
docstring_summary: Construct MACC surrogate model.
docstring_tokens: Construct MACC surrogate model .
function: 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)
function_tokens: def construct_model ( ) : 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 )