remove drop_prob

inference_fuxi.py

@@ -13,7 +13,6 @@ ort.set_default_logger_severity(3)
 parser = argparse.ArgumentParser()
 parser.add_argument('--model', type=str, required=True, help="FuXi onnx model dir")
 parser.add_argument('--input', type=str, required=True, help="The input data file, store in netcdf format")
-parser.add_argument('--drop_prob', type=float, help="dropout prob", default=0)
 parser.add_argument('--input_type', type=str, help="The input type", default="ERA5")
 parser.add_argument('--save_dir', type=str, default="")
 parser.add_argument('--num_steps', type=int, nargs="+", default=[20])
@@ -70,8 +69,6 @@ def run_inference(model_dir, data, num_steps, save_dir=""):
     assert data.lat.values[-1] == -90
 
     input = data.values[None]
-    prob = np.array([args.drop_prob], dtype=np.float32)
-
     print(f'input: {input.shape}, {input.min():.2f} ~ {input.max():.2f}')
     print(f'tembs: {tembs.shape}, {tembs.mean():.4f}')
 
@@ -87,12 +84,15 @@ def run_inference(model_dir, data, num_steps, save_dir=""):
         load_time = time.perf_counter() - start
         print(f'Load model take {load_time:.2f} sec')
 
+        for k, input_name in enumerate(session.get_inputs()):
+            print(k, input_name)
+
         print(f'Inference {stage} ...')
         start = time.perf_counter()
 
         for _ in range(0, num_step):
             temb = tembs[step]
-            new_input, = session.run(None, {'input': input, 'temb': temb, 'prob': prob})
+            new_input, = session.run(None, {'input': input, 'temb': temb})
             output = new_input[:, -1] 
             save_like(output, data, step, save_dir, input_type=args.input_type)
             print(f'stage: {i}, step: {step+1:02d}, output: {output.min():.2f} {output.max():.2f}')