implement warm restart for random forest models (x 2)

Author: ablaom

src/MLJDecisionTreeInterface.jl

@@ -129,13 +129,66 @@ function MMI.fit(
                              m.min_samples_split,
                              m.min_purity_increase;
                              rng=m.rng)
-    cache  = nothing
+    cache  = deepcopy(m)
 
     report = (features=features,)
 
     return (forest, classes_seen, integers_seen), cache, report
 end
 
+function MMI.update(
+    model::RandomForestClassifier,
+    verbosity::Int,
+    old_fitresult,
+    old_model,
+    Xmatrix,
+    yplain,
+    features,
+    classes,
+    )
+
+    only_iterations_have_changed = MMI.is_same_except(model, old_model, :n_trees)
+
+    if !only_iterations_have_changed
+        return MMI.fit(
+            model,
+            verbosity,
+            Xmatrix,
+            yplain,
+            features,
+            classes,
+        )
+    end
+
+    old_forest = old_fitresult[1]
+    Δn_trees = model.n_trees - old_model.n_trees
+    # if `n_trees` drops, then tuncate, otherwise compute more trees
+    if Δn_trees < 0
+        verbosity > 0 && @info "Dropping $(-Δn_trees) trees from the forest. "
+        forest = old_forest[1:model.n_trees]
+    else
+        verbosity > 0 && @info "Adding $Δn_trees trees to the forest. "
+        forest = DT.build_forest(
+            old_forest,
+            yplain, Xmatrix,
+            model.n_subfeatures,
+            model.n_trees,
+            model.sampling_fraction,
+            model.max_depth,
+            model.min_samples_leaf,
+            model.min_samples_split,
+            model.min_purity_increase;
+            rng=model.rng,
+        )
+    end
+
+    fitresult = (forest, old_fitresult[2:3]...)
+    cache = deepcopy(model)
+    report = (features=features,)
+    return fitresult, cache, report
+
+end
+
 MMI.fitted_params(::RandomForestClassifier, (forest,_)) = (forest=forest,)
 
 function MMI.predict(m::RandomForestClassifier, fitresult, Xnew)
@@ -145,7 +198,7 @@ function MMI.predict(m::RandomForestClassifier, fitresult, Xnew)
 end
 
 MMI.reports_feature_importances(::Type{<:RandomForestClassifier}) = true
-
+MMI.iteration_parameter(::Type{<:RandomForestClassifier}) = :n_trees
 
 # # ADA BOOST STUMP CLASSIFIER
 
@@ -266,17 +319,71 @@ function MMI.fit(m::RandomForestRegressor, verbosity::Int, Xmatrix, y, features)
         rng=m.rng
     )
 
-    cache  = nothing
+    cache  = deepcopy(m)
+    report = (features=features,)
+
+    return forest, cache, report
+end
+
+function MMI.update(
+    model::RandomForestRegressor,
+    verbosity::Int,
+    old_forest,
+    old_model,
+    Xmatrix,
+    y,
+    features,
+    )
+
+    only_iterations_have_changed = MMI.is_same_except(model, old_model, :n_trees)
+
+    if !only_iterations_have_changed
+        return MMI.fit(
+            model,
+            verbosity,
+            Xmatrix,
+            y,
+            features,
+        )
+    end
+
+    Δn_trees = model.n_trees - old_model.n_trees
+
+    # if `n_trees` drops, then tuncate, otherwise compute more trees
+    if Δn_trees < 0
+        verbosity > 0 && @info "Dropping $(-Δn_trees) trees from the forest. "
+        forest = old_forest[1:model.n_trees]
+    else
+        verbosity > 0 && @info "Adding $Δn_trees trees to the forest. "
+        forest = DT.build_forest(
+            old_forest,
+            y,
+            Xmatrix,
+            model.n_subfeatures,
+            model.n_trees,
+            model.sampling_fraction,
+            model.max_depth,
+            model.min_samples_leaf,
+            model.min_samples_split,
+            model.min_purity_increase;
+            rng=model.rng
+        )
+    end
+
+    cache = deepcopy(model)
     report = (features=features,)
 
     return forest, cache, report
+
 end
 
 MMI.fitted_params(::RandomForestRegressor, forest) = (forest=forest,)
 
 MMI.predict(::RandomForestRegressor, forest, Xnew) = DT.apply_forest(forest, Xnew)
 
 MMI.reports_feature_importances(::Type{<:RandomForestRegressor}) = true
+MMI.iteration_parameter(::Type{<:RandomForestRegressor}) = :n_trees
+
 
 # # ALIASES FOR TYPE UNIONS
 

test/runtests.jl

@@ -271,4 +271,32 @@ end
     end
 end
 
+@testset "warm restart" begin
+    for M in [:RandomForestClassifier, :RandomForestRegressor]
+        data = (M == :RandomForestClassifier ? make_blobs() : make_regression())
+        quote
+            model = $M(n_trees=4, rng = stable_rng()) # model with 4 trees
+            @test MLJBase.iteration_parameter(model) === :n_trees
+            mach =  machine(model, $data...)
+            fit!(mach, verbosity=0)
+            forest1_4 = fitted_params(mach).forest
+            @test length(forest1_4) ==4
+            mach.model = $M(n_trees=7, rng = stable_rng())
+            @test_logs(
+                (:info, r""),
+                (:info, r"Adding 3 trees"),
+                fit!(mach, verbosity=1),
+            )
+            mach.model = $M(n_trees=5, rng = stable_rng())
+            @test_logs(
+                (:info, r""),
+                (:info, r"Dropping 2 trees"),
+                fit!(mach, verbosity=1),
+            )
+            forest1_5 = fitted_params(mach).forest
+            @test length(forest1_5) == 5
+        end |> eval
+    end
+end
+
 true