Add unload method to module (#12984)

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.

Reviewed By: shoumikhin

Differential Revision: D79184972

Author: silverguo

extension/module/module.h

@@ -194,6 +194,17 @@ 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 True if the method is unloaded, false if no-op.
+   */
+  inline bool unload_method(const std::string& method_name) {
+    return methods_.erase(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 +239,15 @@ class Module {
     return load_forward(nullptr, event_tracer);
   }
 
+  /**
+   * Unload the 'forward' method from the program.
+   *
+   * @returns True if the 'forward' method is unloaded, false if no-op.
+   */
+  inline bool unload_forward() {
+    return unload_method("forward");
+  }
+
   /**
    * Checks if a specific method is loaded.
    *

extension/module/test/module_test.cpp

@@ -91,6 +91,25 @@ 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
+  EXPECT_TRUE(module.unload_method("forward"));
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+  // Try unload method again
+  EXPECT_FALSE(module.unload_method("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_);