Add entrypoints for training with PTE-only

As title.

Author: georgehong

extension/android/executorch_android/android_test_setup.sh

@@ -19,6 +19,8 @@ prepare_add() {
 }
 
 prepare_xor() {
+  python3 -m extension.training.examples.XOR.export_model  --outdir "${BASEDIR}/src/androidTest/resources/"
+  mv "${BASEDIR}/src/androidTest/resources/xor.pte" "${BASEDIR}/src/androidTest/resources/xor_only.pte"
   python3 -m extension.training.examples.XOR.export_model  --outdir "${BASEDIR}/src/androidTest/resources/" --external
 }
 

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/TrainingModuleE2ETest.kt

@@ -108,6 +108,78 @@ class TrainingModuleE2ETest {
         Assert.assertTrue(finalLoss < 0.1f)
     }
 
+    @Test
+    @Throws(IOException::class, URISyntaxException::class)
+    fun testTrainXOR_PTEOnly() {
+        val pteFilePath = "/xor_only.pte"
+
+        val pteFile = File(getTestFilePath(pteFilePath))
+        val pteInputStream = javaClass.getResourceAsStream(pteFilePath)
+        FileUtils.copyInputStreamToFile(pteInputStream, pteFile)
+        pteInputStream.close()
+
+        val module = TrainingModule.load(getTestFilePath(pteFilePath));
+        val params = module.namedParameters("forward")
+
+        Assert.assertEquals(4, params.size)
+        assertContains(params, LIN_WEIGHT)
+        assertContains(params, LIN_BIAS)
+        assertContains(params, LIN2_WEIGHT)
+        assertContains(params, LIN2_BIAS)
+
+        val sgd = SGD.create(params, 0.5);
+        val dataset = listOf<Tensor>(
+            Tensor.fromBlob(floatArrayOf(1.0f, 1.0f), longArrayOf(1, 2)),
+            Tensor.fromBlob(longArrayOf(0), longArrayOf(1)),
+            Tensor.fromBlob(floatArrayOf(0.0f, 0.0f), longArrayOf(1, 2)),
+            Tensor.fromBlob(longArrayOf(0), longArrayOf(1)),
+            Tensor.fromBlob(floatArrayOf(1.0f, 0.0f), longArrayOf(1, 2)),
+            Tensor.fromBlob(longArrayOf(1), longArrayOf(1)),
+            Tensor.fromBlob(floatArrayOf(0.0f, 1.0f), longArrayOf(1, 2)),
+            Tensor.fromBlob(longArrayOf(1), longArrayOf(1)),
+        )
+
+        val numEpochs = 5000;
+        var finalLoss = Float.MAX_VALUE
+
+        for (i in 0 until numEpochs) {
+            val inputDex = 2 * Random.nextInt(dataset.size / 2)
+            val targetDex = inputDex + 1
+            val input = dataset.get(inputDex)
+            val target = dataset.get(targetDex)
+            val out = module.executeForwardBackward("forward", EValue.from(input), EValue.from(target))
+            val gradients = module.namedGradients("forward")
+
+            if (i == 0) {
+                Assert.assertEquals(4, gradients.size)
+                assertContains(gradients, LIN_WEIGHT)
+                assertContains(gradients, LIN_BIAS)
+                assertContains(gradients, LIN2_WEIGHT)
+                assertContains(gradients, LIN2_BIAS)
+            }
+
+            if (i % 500 == 0 || i == numEpochs - 1) {
+                Log.i(
+                    "testTrainXOR_PTEOnly",
+                    String.format(
+                        "Step %d, Loss %f, Input [%.0f, %.0f], Prediction %d, Label %d",
+                        i,
+                        out[0].toTensor().getDataAsFloatArray()[0],
+                        input.getDataAsFloatArray()[0],
+                        input.getDataAsFloatArray()[1],
+                        out[1].toTensor().getDataAsLongArray()[0],
+                        target.getDataAsLongArray()[0]));
+            }
+
+            sgd.step(gradients)
+
+            if (i == numEpochs - 1) {
+                finalLoss = out[0].toTensor().dataAsFloatArray[0]
+            }
+        }
+        Assert.assertTrue(finalLoss < 0.1f)
+    }
+
     companion object {
         private const val LIN_WEIGHT = "net.linear.weight"
         private const val LIN_BIAS = "net.linear.bias"

extension/android/executorch_android/src/main/java/org/pytorch/executorch/TrainingModule.java

@@ -44,9 +44,10 @@ private TrainingModule(String moduleAbsolutePath, String dataAbsolutePath) {
   }
 
   /**
-   * Loads a serialized ExecuTorch module from the specified path on the disk.
+   * Loads a serialized ExecuTorch training module from the specified path on the disk.
    *
    * @param modelPath path to file that contains the serialized ExecuTorch module.
+   * @param dataPath path to file that contains the ExecuTorch module external weights.
    * @return new {@link org.pytorch.executorch.TrainingModule} object which owns the model module.
    */
   public static TrainingModule load(final String modelPath, final String dataPath) {
@@ -61,6 +62,21 @@ public static TrainingModule load(final String modelPath, final String dataPath)
     return new TrainingModule(modelPath, dataPath);
   }
 
+  /**
+   * Loads a serialized ExecuTorch training module from the specified path on the disk.
+   *
+   * @param modelPath path to file that contains the serialized ExecuTorch module. This PTE does not
+   *     rely on external weights.
+   * @return new {@link org.pytorch.executorch.TrainingModule} object which owns the model module.
+   */
+  public static TrainingModule load(final String modelPath) {
+    File modelFile = new File(modelPath);
+    if (!modelFile.canRead() || !modelFile.isFile()) {
+      throw new RuntimeException("Cannot load model path!! " + modelPath);
+    }
+    return new TrainingModule(modelPath, "");
+  }
+
   /**
    * Runs the specified method of this module with the specified arguments.
    *

extension/android/jni/jni_layer_training.cpp

@@ -73,13 +73,16 @@ class ExecuTorchTrainingJni
       facebook::jni::alias_ref<jstring> modelPath,
       facebook::jni::alias_ref<jstring> dataPath) {
     auto modelLoader = FileDataLoader::from(modelPath->toStdString().c_str());
-    auto dataLoader = FileDataLoader::from(dataPath->toStdString().c_str());
+    auto stdStringDataPath = dataPath->toStdString();
     module_ = std::make_unique<training::TrainingModule>(
         std::make_unique<FileDataLoader>(std::move(modelLoader.get())),
         nullptr,
         nullptr,
         nullptr,
-        std::make_unique<FileDataLoader>(std::move(dataLoader.get())));
+        stdStringDataPath.empty()
+            ? nullptr
+            : std::make_unique<FileDataLoader>(std::move(
+                  FileDataLoader::from(stdStringDataPath.c_str()).get())));
   }
 
   static facebook::jni::local_ref<jhybriddata> initHybrid(