From 7243b80a8e5e59ccf138a648658a03346d3d69e2 Mon Sep 17 00:00:00 2001
From: BowenBao <bowbao@microsoft.com>
Date: Mon, 11 May 2020 18:18:57 -0700
Subject: [PATCH] Fix state_dict/load_state_dict with frozen_weight

---
 .../python/onnxruntime_test_ort_trainer.py    |  51 +++++++++++++-
 onnxruntime/test/testdata/ckpt_mnist.pt       | Bin 1595972 -> 1595944 bytes
 orttraining/orttraining/python/ort_trainer.py |  64 ++++++++++++------
 3 files changed, 93 insertions(+), 22 deletions(-)

diff --git a/onnxruntime/test/python/onnxruntime_test_ort_trainer.py b/onnxruntime/test/python/onnxruntime_test_ort_trainer.py
index c1ac2faad1c2d..f55cc4ad2d4aa 100644
--- a/onnxruntime/test/python/onnxruntime_test_ort_trainer.py
+++ b/onnxruntime/test/python/onnxruntime_test_ort_trainer.py
@@ -6,6 +6,7 @@
 import pytest
 import sys
 import copy
+import numpy as np
 from numpy.testing import assert_allclose, assert_array_equal
 
 import onnx
@@ -378,7 +379,7 @@ def testMNISTStateDict(self):
         loss, _ = trainer.train_step(data, target, torch.tensor([learningRate]))
 
         state_dict = trainer.state_dict()
-        assert state_dict.keys() == {'model_.fc1.bias', 'model_.fc1.weight', 'model_.fc2.bias', 'model_.fc2.weight'}
+        assert state_dict.keys() == {'fc1.bias', 'fc1.weight', 'fc2.bias', 'fc2.weight'}
 
     def testMNISTSaveAsONNX(self):
         torch.manual_seed(1)
@@ -446,7 +447,53 @@ def testMNISTFrozenWeight(self):
         data = data.reshape(data.shape[0], -1)
 
         loss, _ = trainer.train_step(data, target, torch.tensor([learningRate]))
-        assert 'fc1.weight' not in trainer.state_dict() and 'fc2.weight' in trainer.state_dict()
+
+        fc1_trainstep_1 = trainer.state_dict()['fc1.weight']
+        fc2_trainstep_1 = trainer.state_dict()['fc2.weight']
+
+        loss, _ = trainer.train_step(data, target, torch.tensor([learningRate]))
+
+        fc1_trainstep_2 = trainer.state_dict()['fc1.weight']
+        fc2_trainstep_2 = trainer.state_dict()['fc2.weight']
+        assert np.array_equal(fc1_trainstep_1, fc1_trainstep_2) and \
+            not np.array_equal(fc2_trainstep_1, fc2_trainstep_2)
+
+    def testMNISTFrozenWeightCheckpoint(self):
+        torch.manual_seed(1)
+        device = torch.device("cuda")
+
+        mnist = MNISTWrapper()
+        train_loader, test_loader = mnist.get_loaders()
+        model, model_desc = mnist.get_model()
+
+        trainer = mnist.get_trainer(model, model_desc, device, frozen_weights=['fc1.weight'])
+
+        learningRate = 0.02
+        epoch = 0
+
+        # do one train step
+        data, target = next(iter(train_loader))
+        data, target = data.to(device), target.to(device)
+        data = data.reshape(data.shape[0], -1)
+
+        loss, _ = trainer.train_step(data, target, torch.tensor([learningRate]))
+
+        # do one eval step
+        data, target = next(iter(train_loader))
+        data, target = data.to(device), target.to(device)
+        data = data.reshape(data.shape[0], -1)
+
+        loss, _ = trainer.eval_step(data, target)
+
+        # save checkpoint, load model and compare
+        state_dict = trainer.state_dict()
+
+        new_model, _ = mnist.get_model()
+        trainer = mnist.get_trainer(new_model, model_desc, device, frozen_weights=['fc1.weight'])
+        trainer.load_state_dict(state_dict)
+
+        ckpt_loss, _ = trainer.eval_step(data, target)
+        assert loss == ckpt_loss
 
     def testBertTrainingBasic(self):
         expected_losses = [
diff --git a/onnxruntime/test/testdata/ckpt_mnist.pt b/onnxruntime/test/testdata/ckpt_mnist.pt
index 01a9a9b4488c1331915e070bc7197a8b0a59b6c7..e3b5723604577474d142c79d9c0632c102f8c1a7 100644
GIT binary patch
delta 382
zcmX@IC2_@;#0l?s%q<NJ4b3eL%`FTpCjN?G;b353m>kQf3zIOgG`2LFJdIIZn+qtD
zmTah3o|>7SQBo+G!K9&)!JHu&!3UJEG=UjpWN2<S`8}h)Cdj0;WFx(#%*5hCMI5Rv
zOiU~$`!T7qfJ~oU$0QB2*WA$D&}i~fChK^p!C;3JGGei<kO@r?<d;HbG(nI}g)FFo
zAWJseGv8%zw#jX`$z=p$CLm@8Viq7~1!6WJW(Q&pAm#*OE+FOxVjdvo1!6uR<_BT{
oAQl8-As`k8Vi6z~1!A%7Ho4*)Z#S1jo)Dhw5^1nmhFilA0MzndXaE2J

delta 388
zcmZ3nCGp6X#0l?s3{6bTjf^Z!ElkWzCjN@x=4W7F$jwhl&575W9LA^%kvFg~H#Rjg
znB32(&Mk<fU>Bn_M1hHkv8ko0(d37W`UXfQq$L~aC1oZS7b<2jX=r3HX9z~{0kvD2
zz-=`!G@ES6q-um@Do9g#YG!&yNuee-ZH7jsCT12ElRKHL<H6=YU02A6Y&cX^ArrD7
z$e%zP3z^Xcfu1j9K^BB~XtOf&UG{A}5*mKfFK|m3Z0GnPenPm3yN#Q%jhm^Bo4Jjf
zrHz}ljhn5Fo4t*jqm7%hjhm~D8>oP%jhnZPo3D+Vzl~d<ja#sdTd0j&xQ$z+ja#&h
LTWlM*ILBK6Etgw`

diff --git a/orttraining/orttraining/python/ort_trainer.py b/orttraining/orttraining/python/ort_trainer.py
index 57b1288dc1f4c..1605688144cf8 100644
--- a/orttraining/orttraining/python/ort_trainer.py
+++ b/orttraining/orttraining/python/ort_trainer.py
@@ -338,8 +338,11 @@ def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, op
         for i, name in enumerate(n.input):
             if name in replace_name_dict:
                 n.input[i] = replace_name_dict[name]
-    assert set([n.name for n in onnx_model.graph.initializer]) == \
-        set([n for n, t in model.model_.named_parameters()]), \
+
+    # onnx model initializer may contain non-trainable registered buffers that are not part
+    # of pytorch model named parameteres.
+    assert set([n for n, t in model.model_.named_parameters()]).issubset(
+        set([n.name for n in onnx_model.graph.initializer])), \
         "Initializer names do not match between PyTorch model and ONNX model, " \
         "please report a bug to ONNX Runtime."
 
@@ -433,8 +436,7 @@ def create_ort_training_session_with_optimizer(model, device, training_optimizer
 
     unused_frozen_weights = [n for n in frozen_weights if n not in [i.name for i in model.graph.initializer]]
     if unused_frozen_weights:
-        warnings.warn("Ignoring {} in frozen_weights as they are not found in model weights."\
-            .format(unused_frozen_weights))
+        raise RuntimeError("{} in frozen_weights not found in model weights.".format(unused_frozen_weights))
 
     weights_to_train = set()
     for initializer in model.graph.initializer:
@@ -763,15 +765,38 @@ def train(self):
     def eval(self):
         self.is_train = False
 
+    def _update_onnx_model_initializers(self, state_tensors):
+        # replace the initializers with new value
+        new_weights = []
+        replace_indices = []
+        i = 0
+        for w in self.onnx_model_.graph.initializer:
+            if w.name in state_tensors:
+                new_weights.append(numpy_helper.from_array(state_tensors[w.name], w.name))
+                replace_indices.append(i)
+            i += 1
+        replace_indices.sort(reverse=True)
+        for w_i in replace_indices:
+            del self.onnx_model_.graph.initializer[w_i]
+        self.onnx_model_.graph.initializer.extend(new_weights)
+
     def state_dict(self):
         if not self.session:
             warnings.warn("ONNXRuntime training session is not initialized yet. "
                           "Please run train_step or eval_step at least once before calling state_dict().")
             return {}
+
+        # extract trained weights
         session_state = self.session.get_state()
         torch_state = {}
         for name in session_state:
             torch_state[name] = torch.from_numpy(session_state[name])
+
+        # extract untrained weights and buffer
+        for n in self.onnx_model_.graph.initializer:
+            if n.name not in torch_state:
+                torch_state[n.name] = torch.from_numpy(numpy_helper.to_array(n))
+
         return torch_state
 
     def load_state_dict(self, state_dict, strict=False):
@@ -783,10 +808,21 @@ def load_state_dict(self, state_dict, strict=False):
             self.strict_ = strict
             return
 
-        session_state = {}
+        # update onnx model from loaded state dict
+        cur_initializers_names = [n.name for n in self.onnx_model_.graph.initializer]
+        new_initializers = {}
+
         for name in state_dict:
-            session_state[name] = state_dict[name].numpy()
-        self.session.load_state(session_state, strict)
+            if name in cur_initializers_names:
+                new_initializers[name] = state_dict[name].numpy()
+            elif strict:
+                raise RuntimeError("Checkpoint tensor: {} is not present in the model.".format(name))
+
+        self._update_onnx_model_initializers(new_initializers)
+
+        # create new session based on updated onnx model
+        self.state_dict_ = None
+        self._init_session()
 
     def save_as_onnx(self, path):
         if not self.session:
@@ -794,19 +830,7 @@ def save_as_onnx(self, path):
                           "Please run train_step or eval_step at least once before calling save_as_onnx().")
             return
         state_tensors = self.session.get_state()
-        # replace the initializers with new value
-        new_weights = []
-        replace_indices = []
-        i = 0
-        for w in self.onnx_model_.graph.initializer:
-            if w.name in state_tensors:
-                new_weights.append(numpy_helper.from_array(state_tensors[w.name], w.name))
-                replace_indices.append(i)
-            i += 1
-        replace_indices.sort(reverse=True)
-        for w_i in replace_indices:
-            del self.onnx_model_.graph.initializer[w_i]
-        self.onnx_model_.graph.initializer.extend(new_weights)
+        self._update_onnx_model_initializers(state_tensors)
 
         with open(path, "wb") as f:
             f.write(self.onnx_model_.SerializeToString())