import hugectr
from mpi4py import MPI
solver = hugectr.CreateSolver(model_name = "dcn",
                              max_eval_batches = 1,
                              batchsize_eval = 16384,
                              batchsize = 16384,
                              lr = 0.001,
                              vvgpu = [[0]],
                              repeat_dataset = True,
                              use_mixed_precision = False,
                              scaler = 1.0,
                              use_cuda_graph = True,
                              metrics_spec = {hugectr.MetricsType.AUC: 1.0})
reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,
                                  source = ["./dcn_data/file_list.txt"],
                                  eval_source = "./dcn_data/file_list_test.txt",
                                  check_type = hugectr.Check_t.Sum,
                                  num_workers = 16)
optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam,
                                    update_type = hugectr.Update_t.Global,
                                    beta1 = 0.9,
                                    beta2 = 0.999,
                                    epsilon = 0.0001)
model = hugectr.Model(solver, reader, optimizer)
model.add(hugectr.Input(label_dim = 1, label_name = "label",
                        dense_dim = 13, dense_name = "dense",
                        data_reader_sparse_param_array = 
                        [hugectr.DataReaderSparseParam("data1", 2, False, 26)]))
model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, 
                            workspace_size_per_gpu_in_mb = 300,
                            embedding_vec_size = 16,
                            combiner = "sum",
                            sparse_embedding_name = "sparse_embedding1",
                            bottom_name = "data1",
                            optimizer = optimizer))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,
                            bottom_names = ["sparse_embedding1"],
                            top_names = ["reshape1"],
                            leading_dim=416))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,
                            bottom_names = ["reshape1", "dense"], top_names = ["concat1"]))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Slice,
                            bottom_names = ["concat1"],
                            top_names = ["slice11", "slice12"],
                            ranges=[(0,429),(0,429)]))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.MultiCross,
                            bottom_names = ["slice11"],
                            top_names = ["multicross1"],
                            num_layers=1))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,
                            bottom_names = ["slice12"],
                            top_names = ["fc1"],
                            num_output=1024))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,
                            bottom_names = ["fc1"],
                            top_names = ["relu1"]))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,
                            bottom_names = ["relu1"],
                            top_names = ["dropout1"],
                            dropout_rate=0.5))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,
                            bottom_names = ["dropout1"],
                            top_names = ["fc2"],
                            num_output=1024))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,
                            bottom_names = ["fc2"],
                            top_names = ["relu2"]))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,
                            bottom_names = ["relu2"],
                            top_names = ["dropout2"],
                            dropout_rate=0.5))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,
                            bottom_names = ["dropout2", "multicross1"],
                            top_names = ["concat2"]))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,
                            bottom_names = ["concat2"],
                            top_names = ["fc3"],
                            num_output=1))
model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.BinaryCrossEntropyLoss,
                            bottom_names = ["fc3", "label"],
                            top_names = ["loss"]))
model.compile()
model.summary()
model.graph_to_json(graph_config_file = "/dump_infer/dcn.json")
model.fit(max_iter = 2300, display = 200, eval_interval = 2000, snapshot = 2000, snapshot_prefix = "/dump_infer/dcn")
model.export_predictions("/dump_infer/dcn_pred_" + str(2000), "/dump_infer/dcn_label_" + str(2000))


from hugectr.inference import InferenceParams, CreateInferenceSession
import numpy as np
batch_size = 16384
num_batches = 1
data_source = "./dcn_data/file_list_test.txt"
inference_params = InferenceParams(model_name = "dcn",
                                max_batchsize = batch_size,
                                hit_rate_threshold = 1.0,
                                dense_model_file = "/dump_infer/dcn_dense_2000.model",
                                sparse_model_files = ["/dump_infer/dcn0_sparse_2000.model"],
                                device_id = 0,
                                use_gpu_embedding_cache = False,
                                cache_size_percentage = 1.0,
                                i64_input_key = False,
                                use_mixed_precision = False,
                                use_cuda_graph = True)
inference_session = CreateInferenceSession("/dump_infer/dcn.json", inference_params)
predictions = inference_session.predict(num_batches = num_batches,
                                      source = data_source,
                                      data_reader_type = hugectr.DataReaderType_t.Norm,
                                      check_type = hugectr.Check_t.Sum)
grount_truth = np.loadtxt("/dump_infer/dcn_pred_2000")
diff = predictions-grount_truth
mse = np.mean(diff*diff)
if mse > 1e-3:
  raise RuntimeError("Too large mse between DCN multi hot inference and training: {}".format(mse))
  sys.exit(1)
else:
  print("DCN multi hot inference results are consistent with those during training, mse: {}".format(mse))
[((3, 9, 13, 76), 'hugectr.CreateSolver', 'hugectr.CreateSolver', (), '', False, 'import hugectr\n'), ((14, 9, 18, 51), 'hugectr.DataReaderParams', 'hugectr.DataReaderParams', (), '', False, 'import hugectr\n'), ((19, 12, 23, 53), 'hugectr.CreateOptimizer', 'hugectr.CreateOptimizer', (), '', False, 'import hugectr\n'), ((24, 8, 24, 48), 'hugectr.Model', 'hugectr.Model', ({(24, 22, 24, 28): 'solver', (24, 30, 24, 36): 'reader', (24, 38, 24, 47): 'optimizer'}, {}), '(solver, reader, optimizer)', False, 'import hugectr\n'), ((94, 19, 104, 54), 'hugectr.inference.InferenceParams', 'InferenceParams', (), '', False, 'from hugectr.inference import InferenceParams, CreateInferenceSession\n'), ((105, 20, 105, 84), 'hugectr.inference.CreateInferenceSession', 'CreateInferenceSession', ({(105, 43, 105, 65): '"""/dump_infer/dcn.json"""', (105, 67, 105, 83): 'inference_params'}, {}), "('/dump_infer/dcn.json', inference_params)", False, 'from hugectr.inference import InferenceParams, CreateInferenceSession\n'), ((110, 15, 110, 54), 'numpy.loadtxt', 'np.loadtxt', ({(110, 26, 110, 53): '"""/dump_infer/dcn_pred_2000"""'}, {}), "('/dump_infer/dcn_pred_2000')", True, 'import numpy as np\n'), ((112, 6, 112, 24), 'numpy.mean', 'np.mean', ({(112, 14, 112, 23): 'diff * diff'}, {}), '(diff * diff)', True, 'import numpy as np\n'), ((29, 10, 35, 50), 'hugectr.SparseEmbedding', 'hugectr.SparseEmbedding', (), '', False, 'import hugectr\n'), ((36, 10, 39, 44), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((40, 10, 41, 90), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((42, 10, 45, 53), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((46, 10, 49, 41), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((50, 10, 53, 44), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((54, 10, 56, 50), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((57, 10, 60, 45), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((61, 10, 64, 44), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((65, 10, 67, 50), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((68, 10, 71, 45), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((72, 10, 74, 52), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((75, 10, 78, 41), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((79, 10, 81, 49), 'hugectr.DenseLayer', 'hugectr.DenseLayer', (), '', False, 'import hugectr\n'), ((28, 25, 28, 77), 'hugectr.DataReaderSparseParam', 'hugectr.DataReaderSparseParam', ({(28, 55, 28, 62): '"""data1"""', (28, 64, 28, 65): '(2)', (28, 67, 28, 72): '(False)', (28, 74, 28, 76): '(26)'}, {}), "('data1', 2, False, 26)", False, 'import hugectr\n')]