[PyTorch Mobile]Move train related files to their own folder (#58205)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/58205

It's worthing moving train related files into their own folder since we are adding more code under the mobile directory.

This diff does that.

Test Plan: run unit tests and ci

Reviewed By: iseeyuan

Differential Revision: D28402432

fbshipit-source-id: cd76a1c4f8ff06508cdc3aad8a169fbf34bb4995

caffe2/CMakeLists.txt

@@ -574,12 +574,12 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
        ${TORCH_SRC_DIR}/csrc/jit/mobile/module.cpp
        ${TORCH_SRC_DIR}/csrc/jit/mobile/observer.cpp
        ${TORCH_SRC_DIR}/csrc/jit/mobile/interpreter.cpp
-       ${TORCH_SRC_DIR}/csrc/jit/mobile/export_data.cpp
        ${TORCH_SRC_DIR}/csrc/jit/mobile/nnc/context.cpp
        ${TORCH_SRC_DIR}/csrc/jit/mobile/nnc/registry.cpp
-       ${TORCH_SRC_DIR}/csrc/jit/mobile/optim/sgd.cpp
-       ${TORCH_SRC_DIR}/csrc/jit/mobile/random.cpp
-       ${TORCH_SRC_DIR}/csrc/jit/mobile/sequential.cpp
+       ${TORCH_SRC_DIR}/csrc/jit/mobile/train/export_data.cpp
+       ${TORCH_SRC_DIR}/csrc/jit/mobile/train/optim/sgd.cpp
+       ${TORCH_SRC_DIR}/csrc/jit/mobile/train/random.cpp
+       ${TORCH_SRC_DIR}/csrc/jit/mobile/train/sequential.cpp
        )
     list(APPEND TORCH_SRCS ${MOBILE_SRCS})
     list(APPEND TORCH_SRCS ${LITE_EAGER_SYMOBLICATION_SRCS})

test/cpp/jit/test_lite_trainer.cpp

@@ -3,13 +3,13 @@
 #include <c10/core/TensorOptions.h>
 #include <torch/csrc/autograd/generated/variable_factories.h>
 #include <torch/csrc/jit/api/module.h>
-#include <torch/csrc/jit/mobile/export_data.h>
 #include <torch/csrc/jit/mobile/import.h>
 #include <torch/csrc/jit/mobile/import_data.h>
 #include <torch/csrc/jit/mobile/module.h>
-#include <torch/csrc/jit/mobile/optim/sgd.h>
-#include <torch/csrc/jit/mobile/random.h>
-#include <torch/csrc/jit/mobile/sequential.h>
+#include <torch/csrc/jit/mobile/train/export_data.h>
+#include <torch/csrc/jit/mobile/train/optim/sgd.h>
+#include <torch/csrc/jit/mobile/train/random.h>
+#include <torch/csrc/jit/mobile/train/sequential.h>
 #include <torch/csrc/jit/serialization/import.h>
 #include <torch/data/dataloader.h>
 #include <torch/torch.h>

tools/build_variables.bzl

@@ -390,7 +390,6 @@ libtorch_extra_sources = libtorch_core_jit_sources + [
     "torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp",
     "torch/csrc/jit/mobile/backport.cpp",
     "torch/csrc/jit/mobile/backport_manager.cpp",
-    "torch/csrc/jit/mobile/export_data.cpp",
     # To be included for eager symbolication in lite interpreter
     # when it is built in libtorch
     "torch/csrc/jit/mobile/debug_info.cpp",
@@ -400,12 +399,13 @@ libtorch_extra_sources = libtorch_core_jit_sources + [
     "torch/csrc/jit/mobile/interpreter.cpp",
     "torch/csrc/jit/mobile/model_compatibility.cpp",
     "torch/csrc/jit/mobile/module.cpp",
-    "torch/csrc/jit/mobile/observer.cpp",
-    "torch/csrc/jit/mobile/optim/sgd.cpp",
-    "torch/csrc/jit/mobile/random.cpp",
-    "torch/csrc/jit/mobile/sequential.cpp",
     "torch/csrc/jit/mobile/nnc/context.cpp",
     "torch/csrc/jit/mobile/nnc/registry.cpp",
+    "torch/csrc/jit/mobile/observer.cpp",
+    "torch/csrc/jit/mobile/train/export_data.cpp",
+    "torch/csrc/jit/mobile/train/optim/sgd.cpp",
+    "torch/csrc/jit/mobile/train/random.cpp",
+    "torch/csrc/jit/mobile/train/sequential.cpp",
     "torch/csrc/jit/serialization/onnx.cpp",
     "torch/csrc/jit/serialization/export.cpp",
     "torch/csrc/jit/serialization/export_module.cpp",

torch/csrc/jit/mobile/train/export_data.cpp

@@ -1,4 +1,4 @@
-#include <torch/csrc/jit/mobile/export_data.h>
+#include <torch/csrc/jit/mobile/train/export_data.h>
 
 #include <torch/csrc/jit/mobile/module.h>
 #include <torch/csrc/jit/runtime/instruction.h>

torch/csrc/jit/mobile/train/optim/sgd.cpp

@@ -1,4 +1,4 @@
-#include <torch/csrc/jit/mobile/optim/sgd.h>
+#include <torch/csrc/jit/mobile/train/optim/sgd.h>
 
 #include <torch/types.h>
 #include <torch/utils.h>

torch/csrc/jit/mobile/train/random.cpp

@@ -1,4 +1,4 @@
-#include <torch/csrc/jit/mobile/random.h>
+#include <torch/csrc/jit/mobile/train/random.h>
 #include <torch/types.h>
 
 #include <algorithm>

torch/csrc/jit/mobile/train/sequential.cpp

@@ -1,4 +1,4 @@
-#include <torch/csrc/jit/mobile/sequential.h>
+#include <torch/csrc/jit/mobile/train/sequential.h>
 #include <torch/types.h>
 
 #include <algorithm>