def main(): # ---------------------------------- # Command-line arguments # ---------------------------------- args = get_args() # ---------------------------------- # Construct layer graph # ---------------------------------- images = lbann.Input(data_field='samples') # Start counting cumulative layers at 1. cumulative_layer_num = 1 log('Input(datum). cumulative_layer_num={n}'.format( n=cumulative_layer_num)) labels = lbann.Input(data_field='labels') cumulative_layer_num += 1 log('Input(labels). cumulative_layer_num={n}'.format( n=cumulative_layer_num)) probs = densenet(args.procs_per_node, 121, cumulative_layer_num, images) # ---------------------------------- # Setup experiment # ---------------------------------- (trainer, model, data_reader_proto, optimizer) = set_up_experiment(args, [images, labels], probs, labels) # ---------------------------------- # Run experiment # ---------------------------------- run_experiment(args, trainer, model, data_reader_proto, optimizer)
def main(): # ---------------------------------- # Command-line arguments # ---------------------------------- args = get_args() # ---------------------------------- # Construct layer graph # ---------------------------------- input_node = lbann.Input() # Start counting cumulative layers at 1. cumulative_layer_num = 1 log('Input. cumulative_layer_num={n}'.format(n=cumulative_layer_num)) (probs, labels) = construct_layer_graph( 121, cumulative_layer_num, input_node) # ---------------------------------- # Setup experiment # ---------------------------------- (model, data_reader_proto, optimizer) = set_up_experiment( args, input_node, probs, labels) # ---------------------------------- # Run experiment # ---------------------------------- # Note: Use `lbann.run` instead for non-LC systems. run_experiment(args, model, data_reader_proto, optimizer)
def _test_o2l_layer_Conv(self, numDims, hasBias): N, C_in, H = (256, 3, 224) C_out = 64 K, P, S, D = (3, 1, 1, 1) G = 1 lbannConv = lbann.Convolution(lbann.Input(), num_dims=numDims, num_output_channels=C_out, has_vectors=False, conv_dims_i=K, conv_pads_i=P, conv_strides_i=S, conv_dilations_i=D, num_groups=G, has_bias=hasBias) inputShapes = {"x": [N, C_in] + [H] * numDims} paramShapes = {"W": [C_out, C_in] + [K] * numDims} if hasBias: paramShapes["b"] = [C_out] node = onnx.helper.make_node("Conv", inputs=["x", "W"] + (["b"] if hasBias else []), outputs=["y"], kernel_shape=[K] * numDims, pads=[P] * (numDims * 2), strides=[S] * numDims, dilations=[D] * numDims, group=G) onnxConv = convertOnnxNode(node, inputShapes, paramShapes).convolution self._assertFields(lbannConv, onnxConv)
def _test_o2l_layer_Pool(self, numDims, poolMode, onnxOp): N, C, H = (256, 3, 224) K, P, S = (3, 1, 1) lbannPooling = lbann.Pooling(lbann.Input(), num_dims=numDims, has_vectors=False, pool_dims_i=K, pool_pads_i=P, pool_strides_i=S, pool_mode=poolMode) inputShapes = {"x": [N, C] + [H] * numDims} node = onnx.helper.make_node( onnxOp, inputs=["x"], outputs=["y"], kernel_shape=[K] * numDims, pads=[P] * (numDims * 2), strides=[S] * numDims, ) onnxPooling = convertOnnxNode(node, inputShapes, {}).pooling self._assertFields(lbannPooling, onnxPooling)
def _test_l2o_layer_convolution(self, numDims, hasBias): N, C_in, H = (256, 3, 224) C_out = 64 K, P, S, D = (3, 1, 1, 1) G = 1 onnxConv = onnx.helper.make_node("Conv", inputs=["x", "W"] + (["b"] if hasBias else []), outputs=["y"], kernel_shape=[K] * numDims, pads=[P] * (numDims * 2), strides=[S] * numDims, dilations=[D] * numDims, group=G) layer = lbann.Convolution(lbann.Input(name="x"), num_dims=numDims, num_output_channels=C_out, has_vectors=False, conv_dims_i=K, conv_pads_i=P, conv_strides_i=S, conv_dilations_i=D, num_groups=G, has_bias=hasBias) lbannConv = parseLbannLayer(layer.export_proto(), {"x_0": (N, C_in, H, H)})["nodes"] self._assertFields(lbannConv, onnxConv)
def main(): # ---------------------------------- # Command-line arguments # ---------------------------------- args = get_args() # ---------------------------------- # Construct layer graph # ---------------------------------- input_node = lbann.Input(target_mode='classification') # Start counting cumulative layers at 1. cumulative_layer_num = 1 log('Input. cumulative_layer_num={n}'.format(n=cumulative_layer_num)) (probs, labels) = construct_layer_graph(args.procs_per_node, 121, cumulative_layer_num, input_node) # ---------------------------------- # Setup experiment # ---------------------------------- (trainer, model, data_reader_proto, optimizer) = set_up_experiment(args, input_node, probs, labels) # ---------------------------------- # Run experiment # ---------------------------------- run_experiment(args, trainer, model, data_reader_proto, optimizer)
def make_model( motif_size, walk_length, num_vertices, embed_dim, learn_rate, num_epochs, embeddings_dir, ): # Layer graph input_ = lbann.Slice( lbann.Input(data_field='samples'), slice_points=str_list([0, motif_size, motif_size+walk_length]), ) motif_indices = lbann.Identity(input_) walk_indices = lbann.Identity(input_) gan = model.gan.CommunityGAN( num_vertices, motif_size, embed_dim, learn_rate, ) loss, real_disc_prob, fake_disc_prob, gen_prob = gan( motif_indices, motif_size, walk_indices, walk_length, ) # Metrics metrics = [ lbann.Metric(real_disc_prob, name='D(real)'), lbann.Metric(fake_disc_prob, name='D(fake)'), lbann.Metric(gen_prob, name='G'), ] # Callbacks callbacks = [ lbann.CallbackPrint(), lbann.CallbackTimer(), lbann.CallbackDumpWeights(directory=embeddings_dir, epoch_interval=num_epochs), ] # Perform computation at double precision for l in lbann.traverse_layer_graph(input_): l.datatype = lbann.DataType.DOUBLE for w in l.weights: w.datatype = lbann.DataType.DOUBLE # Contruct model return lbann.Model( num_epochs, layers=lbann.traverse_layer_graph(input_), objective_function=loss, metrics=metrics, callbacks=callbacks, )
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 test_o2l_layer_Relu(self): N, C, H, W = (100, 200, 300, 400) lbannRelu = lbann.Relu(lbann.Input(), ) node = onnx.helper.make_node( "Relu", inputs=["x"], outputs=["y"], ) onnxRelu = convertOnnxNode(node, {"x": [N, C, H, W]}, {}).relu self._assertFields(lbannRelu, onnxRelu)
def construct_model(): """Model description """ import lbann import lbann.modules fc = lbann.modules.FullyConnectedModule conv = lbann.modules.Convolution2dModule conv1 = conv(20, 3, stride=1, padding=1, name='conv1') conv2 = conv(20, 3, stride=1, padding=1, name='conv2') fc1 = fc(100, name='fc1') fc2 = fc(20, name='fc2') fc3 = fc(num_classes, name='fc3') # Layer graph input = lbann.Input(name='inp_tensor', target_mode='classification') inp_slice = lbann.Slice(input, axis=0, slice_points=str_list([0, dims - 1, dims]), name='inp_slice') xdata = lbann.Identity(inp_slice) ylabel = lbann.Identity(inp_slice, name='gt_y') #NHWC to NCHW x = lbann.Reshape(xdata, dims='14 13 13') x = conv2(conv1(x)) x = lbann.Reshape(x, dims='3380') x = lbann.Dropout(lbann.Relu(fc1(x)), keep_prob=0.5) x = lbann.Dropout(fc2(x), keep_prob=0.5) pred = lbann.Softmax(fc3(x)) gt_label = lbann.OneHot(ylabel, size=num_classes) loss = lbann.CrossEntropy([pred, gt_label], name='loss') acc = lbann.CategoricalAccuracy([pred, gt_label]) layers = list(lbann.traverse_layer_graph(input)) # Setup objective function weights = set() for l in layers: weights.update(l.weights) obj = lbann.ObjectiveFunction(loss) callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()] # Construct model num_epochs = 10 return lbann.Model(num_epochs, weights=weights, layers=layers, metrics=[lbann.Metric(acc, name='accuracy', unit='%')], objective_function=obj, callbacks=callbacks)
def test_l2o_layer_relu(self): N, C, H, W = (100, 200, 300, 400) onnxRelu = onnx.helper.make_node( "Relu", inputs=["x"], outputs=["y"], ) layer = lbann.Relu(lbann.Input(name="x"), ) lbannRelu = parseLbannLayer(layer.export_proto(), {"x_0": (N, C, H, W)})["nodes"] self._assertFields(lbannRelu, onnxRelu)
def _test_o2l_layer_Gemm(self, hasBias): M, N, K = (100, 200, 300) lbannFC = lbann.FullyConnected(lbann.Input(), num_neurons=N, has_bias=hasBias) inputShapes = {"x": [M, K]} paramShapes = {"W": [N, K]} if hasBias: paramShapes["b"] = [N] node = onnx.helper.make_node("Gemm", inputs=["x", "W"] + (["b"] if hasBias else []), outputs=["y"], transB=1) onnxFC = convertOnnxNode(node, inputShapes, paramShapes).fully_connected self._assertFields(lbannFC, onnxFC)
def test_l2o_layer_batch_normalization(self): N, C, H, W = (100, 200, 300, 400) decay = 0.95 epsilon = 1e-6 onnxBN = onnx.helper.make_node( "BatchNormalization", inputs=["x", "scale", "B", "mean", "var"], outputs=["y"], epsilon=epsilon, momentum=decay, spatial=1) layer = lbann.BatchNormalization( lbann.Input(name="x"), decay=decay, epsilon=epsilon, ) lbannBN = parseLbannLayer(layer.export_proto(), {"x_0": (N, C, H, W)})["nodes"] self._assertFields(lbannBN, onnxBN)
def test_o2l_layer_BatchNormalization(self): N, C, H, W = (100, 200, 300, 400) decay = 0.95 epsilon = 1e-6 lbannBN = lbann.BatchNormalization( lbann.Input(), decay=decay, epsilon=epsilon, ) inputShapes = {"x": [N, C, H, W]} paramShapes = {"scale": [C], "B": [C], "mean": [C], "var": [C]} node = onnx.helper.make_node("BatchNormalization", inputs=["x", "scale", "B", "mean", "var"], outputs=["y"], epsilon=epsilon, momentum=decay) onnxBN = convertOnnxNode(node, inputShapes, paramShapes).batch_normalization self._assertFields(lbannBN, onnxBN)
def _test_l2o_layer_pooling(self, numDims, poolMode, onnxOp): N, C, H = (256, 3, 224) K, P, S = (3, 1, 1) onnxPooling = onnx.helper.make_node( onnxOp, inputs=["x"], outputs=["y"], kernel_shape=[K] * numDims, pads=[P] * (numDims * 2), strides=[S] * numDims, ) layer = lbann.Pooling(lbann.Input(name="x"), num_dims=numDims, has_vectors=False, pool_dims_i=K, pool_pads_i=P, pool_strides_i=S, pool_mode=poolMode) lbannPooling = parseLbannLayer(layer.export_proto(), {"x_0": (N, C, H, H)})["nodes"] self._assertFields(lbannPooling, onnxPooling)
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, )
epoch_size=epoch_size, walk_length=walk_length, return_param=return_param, inout_param=inout_param, num_negative_samples=num_negative_samples, ) # ---------------------------------- # Construct layer graph # ---------------------------------- obj = [] metrics = [] # Embedding vectors, including negative sampling # Note: Input is sequence of vertex IDs input_ = lbann.Input(data_field='samples') if args.embeddings == 'distributed': embeddings_weights = lbann.Weights( initializer=lbann.NormalInitializer( mean=0, standard_deviation=1 / args.latent_dim, ), name='embeddings', ) embeddings = lbann.DistEmbedding( input_, weights=embeddings_weights, num_embeddings=num_vertices, embedding_dim=args.latent_dim, sparse_sgd=True, learning_rate=args.learning_rate,
default=100, type=int, help='number of epochs (default: 100)', metavar='NUM') #Add reader prototext lbann.contrib.args.add_optimizer_arguments(parser) args = parser.parse_args() # 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")
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)
action='store_true', help='Keep error signals (default: False)') parser.add_argument('--batch-job', action='store_true', help='Run as a batch job (default: false)') lbann.contrib.args.add_optimizer_arguments( parser, default_optimizer="adam", default_learning_rate=0.001, ) args = parser.parse_args() # Construct layer graph input = lbann.Input(io_buffer='partitioned', target_mode='regression') universes = lbann.Split(input) secrets = lbann.Split(input) 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(input)) # Set parallel_strategy parallel_strategy = get_parallel_strategy_args( sample_groups=args.mini_batch_size, depth_groups=args.depth_groups)
action='store_true', help='Allow LBANN to reclaim error signals buffers (default: False)') parser.add_argument('--batch-job', action='store_true', help='Run as a batch job (default: false)') lbann.contrib.args.add_optimizer_arguments( 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
def construct_model(run_args): """Construct LBANN model. Initial model for ATOM molecular VAE """ import lbann pad_index = run_args.pad_index assert pad_index is not None sequence_length = run_args.sequence_length assert sequence_length is not None, 'should be training seq len + bos + eos' print("sequence length is {}, which is training sequence len + bos + eos".format(sequence_length)) data_layout = "data_parallel" # Layer graph input_ = lbann.Input(data_field='samples',name='inp_data') #Note input assumes to come from encoder script concatenation of input smiles + z inp_slice = lbann.Slice(input_, axis=0, slice_points=str_list([0, sequence_length, sequence_length+run_args.z_dim]), name='inp_slice') inp_smile = lbann.Identity(inp_slice,name='inp_smile') z = lbann.Identity(inp_slice, name='z') wae_loss= [] input_feature_dims = sequence_length embedding_size = run_args.embedding_dim dictionary_size = run_args.num_embeddings assert embedding_size is not None assert dictionary_size is not None save_output = True if run_args.dump_outputs_dir else False print("save output? ", save_output, "out dir ", run_args.dump_outputs_dir) #uncomment below for random sampling #z = lbann.Gaussian(mean=0.0,stdev=1.0, neuron_dims=str(run_args.z_dim)) x = lbann.Slice(inp_smile, slice_points=str_list([0, input_feature_dims])) x = lbann.Identity(x) waemodel = molwae.MolWAE(input_feature_dims, dictionary_size, embedding_size, pad_index,run_args.z_dim,save_output=save_output) x_emb = lbann.Embedding( x, num_embeddings=waemodel.dictionary_size, embedding_dim=waemodel.embedding_size, name='emb', weights=waemodel.emb_weights ) pred, arg_max = waemodel.forward_decoder(x_emb,z) recon = waemodel.compute_loss(x, pred) wae_loss.append(recon) layers = list(lbann.traverse_layer_graph(input_)) # Setup objective function weights = set() for l in layers: weights.update(l.weights) #l2_reg = lbann.L2WeightRegularization(weights=weights, scale=1e-4) #wae_loss.append(l2_reg) print("LEN wae loss ", len(wae_loss)) obj = lbann.ObjectiveFunction(wae_loss) # Initialize check metric callback metrics = [lbann.Metric(recon, name='recon')] callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()] #Dump output (activation) for post processing pred_tensor = lbann.Concatenation(arg_max, name='pred_tensor') conc_out = lbann.Concatenation([input_,pred_tensor], name='conc_out') callbacks.append(lbann.CallbackDumpOutputs(batch_interval=run_args.dump_outputs_interval, execution_modes='test', directory=run_args.dump_outputs_dir, layers=f'{conc_out.name}')) # Construct model return lbann.Model(run_args.num_epochs, weights=weights, layers=layers, objective_function=obj, metrics=metrics, callbacks=callbacks)
def construct_model(run_args): """Construct LBANN model. Initial model for ATOM molecular VAE """ import lbann print("Dump model dir ", run_args.dump_model_dir) assert run_args.dump_model_dir, "evaluate script asssumes a pretrained WAE model" pad_index = run_args.pad_index assert pad_index is not None sequence_length = run_args.sequence_length assert sequence_length is not None print("sequence length is {}".format(sequence_length)) data_layout = "data_parallel" # Layer graph input_ = lbann.Identity(lbann.Input(name='inp', data_field='samples'), name='inp1') wae_loss = [] input_feature_dims = sequence_length embedding_size = run_args.embedding_dim dictionary_size = run_args.num_embeddings assert embedding_size is not None assert dictionary_size is not None save_output = True if run_args.dump_outputs_dir else False print("save output? ", save_output, "out dir ", run_args.dump_outputs_dir) z = lbann.Gaussian(mean=0.0, stdev=1.0, neuron_dims=run_args.z_dim) waemodel = molwae.MolWAE(input_feature_dims, dictionary_size, embedding_size, pad_index, run_args.z_dim, save_output) recon, d1_real, d1_fake, d_adv, arg_max = waemodel(input_, z) zero = lbann.Constant(value=0.0, num_neurons='1', name='zero') one = lbann.Constant(value=1.0, num_neurons='1', name='one') 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') wae_loss.append(recon) 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_weights = [ w for w in weights if not isinstance(w.optimizer, lbann.NoOptimizer) ] l2_reg = lbann.L2WeightRegularization(weights=l2_weights, scale=1e-4) wae_loss.append(d1_real_bce) wae_loss.append(d_adv_bce) wae_loss.append(d1_fake_bce) wae_loss.append(l2_reg) print("LEN wae loss ", len(wae_loss)) obj = lbann.ObjectiveFunction(wae_loss) # Initialize check metric callback metrics = [ lbann.Metric(d_adv_bce, name='adv_loss'), lbann.Metric(recon, name='recon') ] callbacks = [ lbann.CallbackPrint(), #lbann.CallbackStepLearningRate(step=10, amt=0.5), lbann.CallbackTimer() ] callbacks.append( lbann.CallbackReplaceWeights(source_layers=list2str(src_layers), destination_layers=list2str(dst_layers), batch_interval=2)) #Dump output (activation) for post processing if (run_args.dump_outputs_dir): pred_tensor = lbann.Concatenation(arg_max, name='pred_tensor') callbacks.append( lbann.CallbackDumpOutputs( batch_interval=run_args.dump_outputs_interval, execution_modes='test', directory=run_args.dump_outputs_dir, layers=f'inp pred_tensor {waemodel.q_mu.name}')) # Construct model return lbann.Model(run_args.num_epochs, weights=weights, layers=layers, objective_function=obj, metrics=metrics, callbacks=callbacks)
def setup(data_reader_file, name='classifier', num_labels=200, mini_batch_size=128, num_epochs=1000, learning_rate=0.1, bn_statistics_group_size=2, fc_data_layout='model_parallel', warmup_epochs=50, learning_rate_drop_interval=50, learning_rate_drop_factor=0.25, checkpoint_interval=None): # Setup input data input = lbann.Input(target_mode = 'classification') images = lbann.Identity(input) labels = lbann.Identity(input) # Classification network head_cnn = modules.ResNet(bn_statistics_group_size=bn_statistics_group_size) class_fc = lbann.modules.FullyConnectedModule(num_labels, activation=lbann.Softmax, name=f'{name}_fc', data_layout=fc_data_layout) x = head_cnn(images) probs = class_fc(x) # Setup objective function cross_entropy = lbann.CrossEntropy([probs, labels]) l2_reg_weights = set() for l in lbann.traverse_layer_graph(input): if type(l) == lbann.Convolution or type(l) == lbann.FullyConnected: l2_reg_weights.update(l.weights) l2_reg = lbann.L2WeightRegularization(weights=l2_reg_weights, scale=0.0002) obj = lbann.ObjectiveFunction([cross_entropy, l2_reg]) # Setup model metrics = [lbann.Metric(lbann.CategoricalAccuracy([probs, labels]), name='accuracy', unit='%')] callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()] if checkpoint_interval: callbacks.append( lbann.CallbackCheckpoint( checkpoint_dir='ckpt', checkpoint_epochs=5 ) ) # Learning rate schedules if warmup_epochs: callbacks.append( lbann.CallbackLinearGrowthLearningRate( target=learning_rate * mini_batch_size / 128, num_epochs=warmup_epochs ) ) if learning_rate_drop_factor: callbacks.append( lbann.CallbackDropFixedLearningRate( drop_epoch=list(range(0, num_epochs, learning_rate_drop_interval)), amt=learning_rate_drop_factor) ) # Construct model model = lbann.Model(num_epochs, layers=lbann.traverse_layer_graph(input), objective_function=obj, metrics=metrics, callbacks=callbacks) # Setup optimizer # opt = lbann.Adam(learn_rate=learning_rate, beta1=0.9, beta2=0.999, eps=1e-8) opt = lbann.SGD(learn_rate=learning_rate, momentum=0.9) # Load data reader from prototext data_reader_proto = lbann.lbann_pb2.LbannPB() with open(data_reader_file, 'r') as f: google.protobuf.text_format.Merge(f.read(), data_reader_proto) data_reader_proto = data_reader_proto.data_reader for reader_proto in data_reader_proto.reader: reader_proto.python.module_dir = os.path.dirname(os.path.realpath(__file__)) # Return experiment objects return model, data_reader_proto, opt
def make_model(num_vertices=None, node_features=None, num_classes=None, dataset=None, kernel_type='GCN', callbacks=None, num_epochs=1): '''Construct a model DAG using one of the Graph Kernels Args: num_vertices (int): Number of vertices of each graph (default: None) node_features (int): Number of features per noded (default: None) num_classes (int): Number of classes as targets (default: None) dataset (str): Preset data set to use. Either a datset parameter has to be supplied or all of num_vertices, node_features, and num_classes have to be supplied. (default: None) kernel_type (str): Graph Kernel to use in model. Expected one of GCN, GIN, Graph, or GatedGraph (deafult: GCN) callbacks (list): Callbacks for the model. If set to None the model description, GPU usage, training_output, and timer is reported. (default: None) num_epochs (int): Number of epochs to run (default: 1) Returns: (lbann Model Object: A model object with the supplied callbacks, dataset presets, and graph kernels. ''' assert num_vertices != dataset #Ensure atleast one of the values is set if dataset is not None: assert num_vertices is None if dataset == 'MNIST': num_vertices = 75 num_classes = 10 node_features = 1 elif dataset == 'PROTEINS': num_vertices = 100 num_classes = 2 node_features = 3 else: raise Exception("Unkown Dataset") assert num_vertices is not None assert num_classes is not None assert node_features is not None #---------------------------------- # Reshape and Slice Input Tensor #---------------------------------- input_ = lbann.Input(target_mode='classification') # Input dimensions should be (num_vertices * node_features + num_vertices^2 + num_classes ) # Input should have atleast two children since the target is classification data = lbann_Graph_Data(input_, num_vertices, node_features, num_classes) feature_matrix = data.x adj_matrix = data.adj target = data.y #---------------------------------- # Perform Graph Convolution #---------------------------------- if kernel_type == 'GIN': x = GINConvLayer(feature_matrix, adj_matrix) elif kernel_type == 'GCN': x = GCNConvLayer(feature_matrix, adj_matrix) elif kernel_type == 'Graph': x = GraphConvLayer(feature_matrix, adj_matrix) elif kernel_type == 'GatedGraph': x = GATConvLayer(feature_matrix, adj_matrix) else: ValueError( 'Invalid Graph kernel specifier "{}" recieved. Expected one of:\ GIN,GCN,Graph or GatedGraph'.format(kernel_type)) out_channel = x.shape[1] #---------------------------------- # Apply Reduction on Node Features #---------------------------------- average_vector = lbann.Constant(value=1 / num_vertices, num_neurons=str_list([1, num_vertices]), name="Average_Vector") x = x.get_mat(out_channel) x = lbann.MatMul(average_vector, x, name="Node_Feature_Reduction") # X is now a vector with output_channel dimensions x = lbann.Reshape(x, dims=str_list([out_channel]), name="Squeeze") x = lbann.FullyConnected(x, num_neurons=64, name="hidden_layer_1") x = lbann.Relu(x, name="hidden_layer_1_activation") x = lbann.FullyConnected(x, num_neurons=num_classes, name="Output_Fully_Connected") #---------------------------------- # Loss Function and Accuracy s #---------------------------------- probs = lbann.Softmax(x, name="Softmax") loss = lbann.CrossEntropy(probs, target, name="Cross_Entropy_Loss") accuracy = lbann.CategoricalAccuracy(probs, target, name="Accuracy") layers = lbann.traverse_layer_graph(input_) if callbacks is None: print_model = lbann.CallbackPrintModelDescription( ) #Prints initial Model after Setup training_output = lbann.CallbackPrint( interval=1, print_global_stat_only=False) #Prints training progress gpu_usage = lbann.CallbackGPUMemoryUsage() timer = lbann.CallbackTimer() callbacks = [print_model, training_output, gpu_usage, timer] else: if isinstance(callbacks, list): callbacks = callbacks metrics = [lbann.Metric(accuracy, name='accuracy', unit="%")] model = lbann.Model(num_epochs, layers=layers, objective_function=loss, metrics=metrics, callbacks=callbacks) return model
desc = ('Train LeNet on MNIST data using LBANN.') parser = argparse.ArgumentParser(description=desc) lbann.contrib.args.add_scheduler_arguments(parser) parser.add_argument('--job-name', action='store', default='lbann_lenet', type=str, help='scheduler job name (default: lbann_lenet)') args = parser.parse_args() # ---------------------------------- # Construct layer graph # ---------------------------------- # Input data input_ = lbann.Input(target_mode='classification') images = lbann.Identity(input_) labels = lbann.Identity(input_) # LeNet x = lbann.Convolution(images, num_dims=2, num_output_channels=6, num_groups=1, conv_dims_i=5, conv_strides_i=1, conv_dilations_i=1, has_bias=True) x = lbann.Relu(x) x = lbann.Pooling(x, num_dims=2,
def construct_model(run_args): """Construct LBANN model. Initial model for ATOM molecular VAE """ import lbann pad_index = run_args.pad_index assert pad_index is not None sequence_length = run_args.sequence_length assert sequence_length is not None print("sequence length is {}".format(sequence_length)) data_layout = "data_parallel" # Layer graph input_ = lbann.Identity(lbann.Input(name='inp', target_mode="N/A"), name='inp1') vae_loss = [] input_feature_dims = sequence_length embedding_size = run_args.embedding_dim dictionary_size = run_args.num_embeddings assert embedding_size is not None assert dictionary_size is not None kl, recon = molvae.MolVAE(input_feature_dims, dictionary_size, embedding_size, pad_index)(input_) vae_loss.append(kl) vae_loss.append(recon) print("LEN vae loss ", len(vae_loss)) layers = list(lbann.traverse_layer_graph(input_)) # Setup objective function weights = set() for l in layers: weights.update(l.weights) l2_reg = lbann.L2WeightRegularization(weights=weights, scale=5e-4) obj = lbann.ObjectiveFunction(vae_loss) # Initialize check metric callback metrics = [ lbann.Metric(kl, name='kl_loss'), lbann.Metric(recon, name='recon') ] callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()] if (run_args.dump_weights_interval > 0): callbacks.append( lbann.CallbackDumpWeights( directory=run_args.dump_weights_dir, epoch_interval=run_args.dump_weights_interval)) if (run_args.ltfb): send_name = ('' if run_args.weights_to_send == 'All' else run_args.weights_to_send) #hack for Merlin empty string weights_to_ex = [w.name for w in weights if send_name in w.name] print("LTFB Weights to exchange ", weights_to_ex) callbacks.append( lbann.CallbackLTFB(batch_interval=run_args.ltfb_batch_interval, metric='recon', weights=list2str(weights_to_ex), low_score_wins=True, exchange_hyperparameters=True)) if (run_args.warmup): callbacks.append( lbann.CallbackLinearGrowthLearningRate(target=run_args.lr / 512 * run_args.batch_size, num_epochs=5)) # Construct model return lbann.Model(run_args.num_epochs, weights=weights, layers=layers, objective_function=obj, metrics=metrics, callbacks=callbacks)
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)
def construct_model(run_args): """Construct LBANN model. Initial model for ATOM molecular VAE """ import lbann print("Dump model dir ", run_args.dump_model_dir) assert run_args.dump_model_dir, "evaluate script asssumes a pretrained WAE model" pad_index = run_args.pad_index assert pad_index is not None sequence_length = run_args.sequence_length assert sequence_length is not None print("sequence length is {}".format(sequence_length)) data_layout = "data_parallel" # Layer graph input_ = lbann.Identity(lbann.Input(name='inp', target_mode="N/A"), name='inp1') wae_loss = [] input_feature_dims = sequence_length embedding_size = run_args.embedding_dim dictionary_size = run_args.num_embeddings assert embedding_size is not None assert dictionary_size is not None save_output = False print("save output? ", save_output, "out dir ", run_args.dump_outputs_dir) z = lbann.Gaussian(mean=0.0, stdev=1.0, neuron_dims="128") x = lbann.Slice(input_, slice_points=str_list([0, input_feature_dims])) x = lbann.Identity(x) waemodel = molwae.MolWAE(input_feature_dims, dictionary_size, embedding_size, pad_index, save_output) x_emb = lbann.Embedding(x, num_embeddings=waemodel.dictionary_size, embedding_dim=waemodel.embedding_size, name='emb', weights=waemodel.emb_weights) latentz = waemodel.forward_encoder(x_emb) fake_loss = lbann.MeanAbsoluteError(latentz, z) layers = list(lbann.traverse_layer_graph(input_)) # Setup objective function weights = set() for l in layers: weights.update(l.weights) obj = lbann.ObjectiveFunction(fake_loss) callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()] #Dump output (activation) for post processing conc_out = lbann.Concatenation([input_, latentz], name='conc_out') callbacks.append( lbann.CallbackDumpOutputs( batch_interval=run_args.dump_outputs_interval, execution_modes='test', directory=run_args.dump_outputs_dir, layers=f'{conc_out.name}')) # Construct model return lbann.Model(run_args.num_epochs, weights=weights, layers=layers, objective_function=obj, callbacks=callbacks)
name='decoder_embeddings', ) # ---------------------------------- # Construct layer graph # ---------------------------------- # Properties of graph and random walk num_graph_nodes = dataset.max_graph_node_id() + 1 walk_length = dataset.walk_context_length num_negative_samples = dataset.num_negative_samples input_size = dataset.sample_dims()[0] # Embedding vectors, including negative sampling # Note: Input is sequence of graph node IDs input_ = lbann.Identity(lbann.Input()) input_slice = lbann.Slice( input_, slice_points=f'0 {num_negative_samples+1} {input_size}' ) decoder_embeddings = lbann.Embedding( input_slice, weights=decoder_embeddings_weights, num_embeddings=num_graph_nodes, embedding_dim=args.latent_dim, ) encoder_embeddings = lbann.Embedding( input_slice, weights=encoder_embeddings_weights, num_embeddings=num_graph_nodes, embedding_dim=args.latent_dim,