Add unload method to module

Summary: We would like to do grid search on device when training the model, ideally we can unload/load the model weights after trained with each specific hyperparameter set.

Differential Revision: D79184972

Author: silverguo

extension/module/module.cpp

@@ -228,6 +228,12 @@ runtime::Error Module::load_method(
   return runtime::Error::Ok;
 }
 
+runtime::Error Module::unload_method(const std::string& method_name) {
+  ET_CHECK_OK_OR_RETURN_ERROR(load());
+  methods_.erase(method_name);
+  return runtime::Error::Ok;
+}
+
 ET_NODISCARD runtime::Result<Method*> Module::method(
     const std::string& method_name) {
   ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));

extension/module/module.h

@@ -194,6 +194,16 @@ class Module {
     return load_method(method_name, nullptr, event_tracer);
   }
 
+  /**
+   * Unload a specific method from the program.
+   *
+   * @param[in] method_name The name of the method to unload.
+   *
+   * @returns An Error to indicate success or failure.
+   */
+  ET_NODISCARD
+  runtime::Error unload_method(const std::string& method_name);
+
   /**
    * Get a method by it's name. Not recommended to use this method directly as
    * an end user. It's exposed to allow for composability of module in apis that
@@ -228,6 +238,16 @@ class Module {
     return load_forward(nullptr, event_tracer);
   }
 
+  /**
+   * Unload the 'forward' method from the program.
+   *
+   * @returns An Error to indicate success or failure.
+   */
+  ET_NODISCARD
+  runtime::Error unload_forward() {
+    return unload_method("forward");
+  };
+
   /**
    * Checks if a specific method is loaded.
    *

extension/module/test/module_test.cpp

@@ -91,6 +91,24 @@ TEST_F(ModuleTest, TestLoadMethod) {
   EXPECT_TRUE(module.is_loaded());
 }
 
+TEST_F(ModuleTest, TestUnloadMethod) {
+  Module module(model_path_);
+
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+  const auto errorLoad = module.load_method("forward");
+  EXPECT_EQ(errorLoad, Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+  // Unload method
+  const auto errorUnload = module.unload_method("forward");
+  EXPECT_EQ(errorUnload, Error::Ok);
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+  // Load method again
+  const auto errorReload = module.load_method("forward");
+  EXPECT_EQ(errorReload, Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+  EXPECT_TRUE(module.is_loaded());
+}
+
 TEST_F(ModuleTest, TestLoadNonExistentMethod) {
   Module module(model_path_);