Batched inference output.

linqs · Apr 26, 2024 · c0c0eeb · c0c0eeb
1 parent 4a62a96
commit c0c0eeb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/psl-java/src/main/java/org/linqs/psl/config/RuntimeOptions.java b/psl-java/src/main/java/org/linqs/psl/config/RuntimeOptions.java
@@ -94,6 +94,13 @@ public class RuntimeOptions {
         "Use the specified InferenceApplication when running inference."
     );
 
+    public static final Option INFERENCE_OUTPUT_BATCHED_RESULTS = new Option(
+        "runtime.inference.output.batched.results",
+        false,
+        "Whether to output the inferred atoms after inference and organize output by batches."
+        + " This is useful if the neural component is batching."
+    );
+
     public static final Option INFERENCE_OUTPUT_RESULTS = new Option(
         "runtime.inference.output.results",
         true,

diff --git a/psl-java/src/main/java/org/linqs/psl/runtime/Runtime.java b/psl-java/src/main/java/org/linqs/psl/runtime/Runtime.java
@@ -431,6 +431,8 @@ protected void runInferenceInternal(RuntimeConfig config, Model model, RuntimeRe
 
         // Run inference.
         boolean runInference = true;
+        int batch = 0;
+        String outputDir = RuntimeOptions.INFERENCE_OUTPUT_RESULTS_DIR.getString();
         while (runInference) {
             DeepPredicate.predictAllDeepPredicates();
 
@@ -439,20 +441,28 @@ protected void runInferenceInternal(RuntimeConfig config, Model model, RuntimeRe
             log.info("Inference complete.");
 
             if (RuntimeOptions.INFERENCE_OUTPUT_RESULTS.getBoolean()) {
-                String outputDir = RuntimeOptions.INFERENCE_OUTPUT_RESULTS_DIR.getString();
                 if (outputDir == null) {
                     log.info("Writing inferred predicates to stdout.");
                     targetDatabase.outputRandomVariableAtoms();
                 } else {
-                    log.info("Writing inferred predicates to directory: " + outputDir);
-                    targetDatabase.outputRandomVariableAtoms(outputDir);
+                    if (RuntimeOptions.INFERENCE_OUTPUT_BATCHED_RESULTS.getBoolean()) {
+                        String batchOutputDir = FileUtils.makePath(outputDir, String.format("batch_%d", batch));
+
+                        log.info("Writing inferred predicates to directory: " + batchOutputDir);
+                        targetDatabase.outputRandomVariableAtoms(batchOutputDir);
+                    } else {
+                        log.info("Writing inferred predicates to directory: " + outputDir);
+                        targetDatabase.outputRandomVariableAtoms(outputDir);
+                    }
                 }
             }
 
             DeepPredicate.evalAllDeepPredicates();
 
             DeepPredicate.nextBatchAllDeepPredicates();
             runInference = !DeepPredicate.isEpochCompleteAllDeepPredicates();
+
+            batch++;
         }
         DeepPredicate.epochEndAllDeepPredicates();