Reshape pretrained weights ULMFiT LM

AdarshKumar712 · AdarshKumar712 · commit 6cd9824188b9 · 2021-05-21T05:08:55.000-07:00
diff --git a/src/ULMFiT/pretrain_lm.jl b/src/ULMFiT/pretrain_lm.jl
@@ -163,6 +163,11 @@ end
 # To load model
 function load_model!(lm::LanguageModel, filepath::String)
     BSON.@load filepath weights
+    # reshape saved weights to match Recurr (h, c) shape
+    layers = [5, 6, 10, 11, 15, 16]
+    for l in layers
+        weights[l] = reshape(weights[l], length(weights[l]), 1)
+    end
     Flux.loadparams!(lm, weights)
 end
 
diff --git a/test/ulmfit.jl b/test/ulmfit.jl
@@ -95,6 +95,12 @@ end
     @test length(ULMFiT.get_trainable_params(lm.layers)) == 10
 
     pretrained_weights = BSON.load(datadep"Pretrained ULMFiT Language Model/ulmfit_lm_en.bson")
+    # reshape weights of (h, c) 
+    layers = [5, 6, 10, 11, 15, 16]
+    for i in layers 
+       pretrained_weights[:weights][i] = reshape(pretrained_weights[:weights][i], length(pretrained_weights[:weights][i]), 1)
+    end
+
     @test length(pretrained_weights[:weights]) == 16
     @test all(size.(params(lm)) .== size.(pretrained_weights[:weights]))
 end
