模型的載入與儲存
模型的載入與儲存
對於模型的載入與儲存，常用的場景有：
• 將已經訓練一段時間的模型儲存，方便下次繼續訓練
• 將訓練好的模型儲存，方便後續直接部署使用
嚴格來說，尚未訓練好的模型的儲存，稱為 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,
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 中。