作業函式的定義與呼叫

在 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。

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 還支援非同步呼叫，具體可以參閱專題獲取作業函式的結果。

作業函式的定義與呼叫

函式的定義與呼叫

C 函式的定義與呼叫

3 1 函式的定義與呼叫

作業函式的定義與呼叫

函式的定義與呼叫

C 函式的定義與呼叫

3 1 函式的定義與呼叫

相關推薦