作業函數的定義與調用
在 OneFlow 中，將訓練、預測任務封裝在一個函數中，統稱為作業函數(job function)，作業函數聯系用戶的業務邏輯與 OneFlow 管理的計算資源。
在 OneFlow 中，被 @oneflow.global_function 裝飾器所修飾的 python 函數，就是 OneFlow 作業函數。
主要在作業函數中定義網絡模型的結構、選擇優化指標；此外，還可以將訓練有關的超參及環境配置當做參數，傳遞給作業函數(如:下面例子中的：get_train_config())，OneFlow 會根據設置為管理內存、GPU 等硬件資源。
本文中將具體學習：
? 如何定義和調用作業函數
? 如何獲取作業函數的返回值
作業函數與 OneFlow 運行流程的關系
作業函數分為定義和調用兩個階段。
這與 OneFlow 本身的運行機制有關，簡化地說，OneFlow Python 層接口，只是在描述網絡模型和訓練環境的配置信息，這些信息將傳遞給底層的 C++ 代碼，經過編譯、優化等工作得到計算圖，最終交給 OneFlow 運行時(runtime)，由 OneFlow 運行時執行。
因為定義作業函數只是做描述工作，在這個階段，并沒有實際的數據，而只能通過規定網絡節點的形狀、數據類型等信息，起到 數據占位符 的作用，方便 OneFlow 的編譯構圖過程進行模型推理。
作業函數的調用，發生在 OneFlow runtime 已經啟動后，可以通過調用作業函數，向其傳遞真實的數據，并獲取返回結果。
以下將具體介紹作業函數的定義與調用方法。
作業函數的定義
將模型封裝在 Python 中，再使用oneflow.global_function修飾符進行修飾。就完成了作業函數的定義。
作業函數主要描述兩方面的事情：
? 模型結構
? 訓練過程中的優化目標
以下代碼示例中，構建了一個 mlp 模型。并且將由 flow.nn.sparse_softmax_cross_entropy_with_logits 計算得到交叉熵損失結果作為優化目標。
@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.Callback[tp.Numpy]:
# mlp
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=“hidden”,
)
logits = flow.layers.dense(
hidden, 10, kernel_initializer=initializer, name=“output”
)

loss = flow.nn.sparse_softmax_cross_entropy_with_logits(labels, logits, name="softmax_loss"
)
lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
return loss

oneflow.global_function 的參數
oneflow.global_function 修飾符接受兩個參數，分別是 type 與 function_config。
? type 參數接受字符串，只能設定為 train 或者 predict，當定義一個訓練模型時，設定為 train，當定義測試模型時，設定為 predict
? function_config 參數接受一個 oneflow.function_config() 所構造的對象，在 function_config 對象中，可以通過成員方法或屬性，進行相關配置。如以下代碼：
def get_train_config():
config = flow.function_config()
config.default_data_type(flow.float)
return config
設置了默認數據類型，然后，可以在向 global_function 裝飾器傳遞這個function_config 對象：
@flow.global_function(type=“train”, function_config=get_train_config())
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:
包含以上代碼的完整示例可見文章Consistent 與 Mirrored 視角中的 mixed_parallel_mlp.py
數據占位符
注意，以上的 images、logits、labels、loss等對象，在定義作業函數時，并沒有實際的數據。作用只是 描述數據的形狀和屬性 ，起到 占位符 的作用。
在作業函數的參數中的數據占位符，使用 oneflow.typing 下的Numpy.Placeholder、ListNumpy.Placeholder、ListListNumpy.Placeholder，注解作業函數的參數，對應作業函數調用時，傳遞 numpy 數據對象。
除了oneflow.typing下的幾種類型外，不出現在參數中，而由 OneFlow 的算子或層產生的變量，如以上代碼中的reshape、hidden、logits、loss等，也都起到了數據占位符的作用。
不管是以上提及的哪種變量，都直接或間接繼承自 OneFlow 的 BlobDef 基類，OneFlow 中把這種對象類型統稱為 Blob。
Blob 在作業函數定義時，均無真實數據，均只起到數據占位方便框架推理的作用。
作業函數的返回值
之所以在上文中強調數據占位符 Blob 的概念，因為作業函數的返回值是不能任意指定的，必須是 Blob 類型的對象，或者存有 Blob 對象的容器。
如以上代碼的中所返回的 loss，它就是 Blob 類型。
作業函數的返回值，需要通過注解聲明，比如以上代碼中的 -> tp.Numpy，表示返回1個 Blob。
再比如，可以通過注解聲明返回值類型為 -> Tuple[tp.Numpy, tp.Numpy]，表示返回1個 tuple，該 tuple 中有2個 Blob 對象。
具體的使用例子，可以參考獲取作業函數的結果
作業函數的調用
OneFlow 利用函數修飾符將普通 Python 函數轉變為 OneFlow 特有的作業函數的過程，對于用戶而言是無感、透明的。
可以像調用普通的 Python 函數一樣調用作業函數。每一次調用，OneFlow 都會在框架內部完成正向傳播、反向傳播、參數更新等一系列事情。
以下代碼，獲取數據之后，向 train_job 作業函數傳遞參數并調用，打印 loss。
(train_images, train_labels), (test_images, test_labels) = flow.data.load_mnist(
BATCH_SIZE
)

for epoch in range(3):for i, (images, labels) in enumerate(zip(train_images, train_labels)):loss = train_job(images, labels)if i % 20 == 0:print(loss.mean())

可以看到，通過調用作業函數 train_job 直接返回了 numpy 數據。
以上展示的調用方式是同步方式， OneFlow 還支持異步調用，具體可以參閱專題獲取作業函數的結果。

總結

以上是生活随笔為你收集整理的作业函数的定义与调用的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。