模型的加载与保存
模型的加載與保存
對于模型的加載與保存，常用的場景有：
? 將已經訓練一段時間的模型保存，方便下次繼續訓練
? 將訓練好的模型保存，方便后續直接部署使用
嚴格來說，尚未訓練好的模型的保存，稱為 checkpoint 或者 snapshot 。與保存已訓練好的模型（model saving） ，在概念上，略有不同。
不過，在 OneFlow 中，無論模型是否訓練完畢，都使用 統一的接口 將其保存，因此，在其它框架中看到的model、checkpoint、snapshot 等表述，在 OneFlow 中不做區分。
在 OneFlow 中，flow.checkpoint 名稱空間下有模保存、加載的接口。
本文將介紹：
? 如何創建模型參數
? 如何保存/加載模型
? OneFlow 模型的存儲結構
? 如何微調與擴展模型
get_variable 創建或獲取參數
可以使用 oneflow.get_variable 方法創造或者獲取一個對象，該對象可以用于在全局作業函數中交互信息；當調用 oneflow.get_all_variables 和 oneflow.load_variables 接口時，可以獲取或更新 get_variable 創建的對象的值。
因為這個特點，get_variable 創建的對象，常用于存儲模型參數。實際上，OneFlow 中很多較高層接口（如 oneflow.layers.conv2d），內部使用 get_variable 創建模型參數。
流程
get_variable 需要一個指定一個 name 參數，該參數作為創建對象的標識。
如果 name 指定的值在當前上下文環境中已經存在，那么 get_variable 會取出已有對象，并返回。
如果 name 指定的值不存在，則 get_varialbe 內部會創建一個 blob 對象，并返回。
使用 get_variable 創建對象
oneflow.get_variable 的原型如下：
def get_variable(
name,
shape=None,
dtype=None,
initializer=None,
regularizer=None,
trainable=None,
model_name=None,
random_seed=None,
distribute=distribute_util.broadcast(),
)
以下是 oneflow.layers.conv2d 中，使用 get_variable 創造參數變量，并進一步構建網絡的例子：
#…
weight = flow.get_variable(
weight_name if weight_name else name_prefix + “-weight”,
shape=weight_shape,
dtype=inputs.dtype,
initializer=kernel_initializer
if kernel_initializer is not None
else flow.constant_initializer(0),
regularizer=kernel_regularizer,
trainable=trainable,
model_name=“weight”,
)
output = flow.nn.conv2d(inputs, weight, strides, padding, data_format, dilation_rate, groups=groups, name=name
)
#...
initializer 設置初始化方式
在上文中已經看到，在調用 get_variable 時，通過設置初始化器 initializer 來指定參數的初始化方式，OneFlow 中提供了多種初始化器，可以在 oneflow 模塊下查看。
在靜態圖機制下，設置 initializer 后，參數初始化工作由 OneFlow 框架自動完成。
OneFlow 目前支持的 initializer 列舉如下，點擊鏈接可以查看相關算法：
? constant_initializer
? zeros_initializer
? ones_initializer
? random_uniform_initializer
? random_normal_initializer
? truncated_normal_initializer
? glorot_uniform_initializer
? glorot_normal_initializer
? variance_scaling_initializer
? kaiming_initializer
? xavier_normal_initializer
? xavier_uniform_initializer
OneFlow 模型的 Python 接口
注意：由于多版本兼容的原因，使用本節介紹的接口，在腳本中都需先配置：
flow.config.enable_legacy_model_io(False)
獲取/更新 variable 對象的值
可以使用以下兩個接口，獲取或更新作業函數中由 oneflow.get_variable 所創建的 variable 對象的值
? oneflow.get_all_variables : 獲取所有作業函數中的的 variable 對象
? oneflow.load_variables : 更新作業函數中的 variable 對象
oneflow.get_all_variables 會返回一個字典，字典的 key 就是創建 variable 時指定的 name，key 對應的 value 就是一個張量對象，該張量對象有 numpy() 轉為 numpy 數組。
比如，在作業函數中創建了名為 myblob 的對象：
@flow.global_function()
def job() -> tp.Numpy:
…
myblob = flow.get_variable(“myblob”,
shape=(3,3),
initializer=flow.random_normal_initializer()
)
…
如果想打印 myblob 的值，可以調用：
…
for epoch in range(20):
…
job()
all_variables = flow.get_all_variables()
print(all_variables[“myblob”].numpy())
…
其中的 flow.get_all_variables 獲取到了字典，all_variables[“myblob”].numpy() 獲取了 myblob 對象并將其轉為 numpy 數組。
與 get_all_variables 相反，可以使用 oneflow.load_variables 更新 varialbe 對象的值。 oneflow.load_variables 的原型如下：
def load_variables(value_dict, ignore_mismatch = True)
使用 load_variables 前，要準備一個字典，該字典的 key 為創建 variable 時指定的 name，value 是 numpy 數組；將字典傳遞給 load_variables 后，load_variables 會將根據 key 找到作業函數中的 variable 對象，并更新值。
如以下代碼：
@flow.global_function(type=“predict”)
def job() -> tp.Numpy:
myblob = flow.get_variable(“myblob”,
shape=(3,3),
initializer=flow.random_normal_initializer()
)
return myblob
myvardict = {“myblob”: np.ones((3,3)).astype(np.float32)}
flow.load_variables(myvardict)
print(flow.get_all_variables()[“myblob”].numpy())
雖然選擇了 random_normal_initializer 的初始化方式，但是因為 flow.load_variables(myvardict) 更新了 myblob 的值，所以最終輸出結果是：
[[1. 1. 1.]
[1. 1. 1.]
[1. 1. 1.]]
模型的保存與加載
通過以下兩個方法，可以保存/加載模型：
? oneflow.checkpoint.save : 負責保存當前的模型到指定路徑
? oneflow.checkpoint.get : 從指定路徑中導入模型
save 的原型如下，可以將模型保存至 path 所指定的路徑。
def save(path, var_dict=None)
可選參數 var_dict 如果不為 None，則將 var_dict 中指定的對象保存到指定路徑。
get 的原型如下，可以加載之前已經保存的，由 path 路徑所指定的模型。
def get(path)
它將返回一個字典，該字典可以用上文介紹的 load_variables 方法更新到模型中：
flow.load_variables(flow.checkpoint.get(save_dir))
注意：
? save 參數所指定路徑對應的目錄要么不存在，要么應該為空目錄，否則 save 會報錯(防止覆蓋掉原有保存的模型)
? OneFlow 模型以一定的組織形式保存在指定的路徑中，具體結構參見下文中的 OneFlow 模型的存儲結構
? 雖然 OneFlow 對 save 的頻率沒有限制，但是過高的保存頻率，會加重磁盤及帶寬等資源的負擔。
OneFlow 模型的存儲結構
OneFlow 模型是一組已經被訓練好的網絡的 參數值 。模型所保存的路徑下，有多個子目錄，每個子目錄對應了 作業函數 中模型的 name。 比如，先通過代碼定義以下的模型：
def lenet(data, train=False):
initializer = flow.truncated_normal(0.1)
conv1 = flow.layers.conv2d(
data,
32,
5,
padding=“SAME”,
activation=flow.nn.relu,
name=“conv1”,
kernel_initializer=initializer,
)
pool1 = flow.nn.max_pool2d(
conv1, ksize=2, strides=2, padding=“SAME”, name=“pool1”, data_format=“NCHW”
)
conv2 = flow.layers.conv2d(
pool1,
64,
5,
padding=“SAME”,
activation=flow.nn.relu,
name=“conv2”,
kernel_initializer=initializer,
)
pool2 = flow.nn.max_pool2d(
conv2, ksize=2, strides=2, padding=“SAME”, name=“pool2”, data_format=“NCHW”
)
reshape = flow.reshape(pool2, [pool2.shape[0], -1])
hidden = flow.layers.dense(
reshape,
512,
activation=flow.nn.relu,
kernel_initializer=initializer,
name=“dense1”,
)
if train:
hidden = flow.nn.dropout(hidden, rate=0.5, name=“dropout”)
return flow.layers.dense(hidden, 10, kernel_initializer=initializer, name=“dense2”)
假設在訓練過程中，調用以下代碼保存模型：
flow.checkpoint.save(’./lenet_models_name’)
那么 lenet_models_name 及其子目錄結構為：
lenet_models_name/
├── conv1-bias
│ ├── meta
│ └── out
├── conv1-weight
│ ├── meta
│ └── out
├── conv2-bias
│ ├── meta
│ └── out
├── conv2-weight
│ ├── meta
│ └── out
├── dense1-bias
│ ├── meta
│ └── out
├── dense1-weight
│ ├── meta
│ └── out
├── dense2-bias
│ ├── meta
│ └── out
├── dense2-weight
│ ├── meta
│ └── out
├── snapshot_done
└── System-Train-TrainStep-train_job
├── meta
└── out
可以看到：
? 作業函數中的網絡模型，每個變量對應一個子目錄
? 以上每個子目錄中，都有 out 和 meta 文件，out 以二進制的形式存儲了網絡參數的值，meta 以文本的形式存儲了網絡的結構信息
? snapshot_done 是一個空文件，如果存在，表示網絡已經訓練完成
? System-Train-TrainStep-train_job 中保存有快照的訓練步數
模型的微調與擴展
在模型的微調和遷移學習中，經常需要：
? 模型中的一部分參數加載自原有模型
? 模型中的另一部分（新增的）參數需要初始化
可以使用 oneflow.load_variables 完成以上操作。以下舉一個用于闡述概念的簡單例子。
首先，先定義一個模型，訓練后保存至 ./mlp_models_1：
@flow.global_function(type=“train”)
def train_job(
images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
with flow.scope.placement(“cpu”, “0:0”):
initializer = flow.truncated_normal(0.1)
reshape = flow.reshape(images, [images.shape[0], -1])
hidden = flow.layers.dense(
reshape,
512,
activation=flow.nn.relu,
kernel_initializer=initializer,
name=“dense1”,
)
dense2 = flow.layers.dense(
hidden, 10, kernel_initializer=initializer, name=“dense2”
)
loss = flow.nn.sparse_softmax_cross_entropy_with_logits(labels, dense2)lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)return loss
然后，拓展網絡結構，為以上模型多增加一層 dense3：
@flow.global_function(type=“train”)
def train_job(
images: tp.Numpy.Placeholder((BATCH_SIZE, 1, 28, 28), dtype=flow.float),
labels: tp.Numpy.Placeholder((BATCH_SIZE,), dtype=flow.int32),
) -> tp.Numpy:
with flow.scope.placement(“cpu”, “0:0”):
#… 原有網絡結構
dense3 = flow.layers.dense(dense2, 10, kernel_initializer=initializer, name="dense3")loss = flow.nn.sparse_softmax_cross_entropy_with_logits(labels, dense3)#...
最后，從原來保存的模型加載參數，并開始訓練：
if name == “main”:
flow.load_variables(flow.checkpoint.get("./mlp_models_1"))
(train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(BATCH_SIZE, BATCH_SIZE
)
for i, (images, labels) in enumerate(zip(train_images, train_labels)):loss = train_job(images, labels)if i % 20 == 0:print(loss.mean())
flow.checkpoint.save("./mlp_ext_models_1")
新增的 dense3 層參數，在原模型中不存在，OneFlow 會自動初始化它們的值。
代碼
腳本 mlp_mnist_origin.py 中構建了“骨干網絡”，并將訓練好的模型保存至 ./mlp_models_1。
運行：
wget https://docs.oneflow.org/code/basics_topics/mlp_mnist_origin.py
python3 mlp_mnist_origin.py
訓練完成后，將會在當前工作路徑下得到 mlp_models_1 目錄。
腳本 mlp_mnist_finetune.py 中的網絡在原有基礎上進行“微調”（為骨干網絡增加一層dense3）后，加載 ./mlp_models_1，并繼續訓練。
運行：
wget https://docs.oneflow.org/code/basics_topics/mlp_mnist_finetune.py
python3 mlp_mnist_finetune.py
微調后的模型，保存在 ./mlp_ext_models_1 中。
總結
