1. 程式人生 > 其它 >TensorFlow 2.0 - TFRecord儲存資料集、@tf.function圖執行模式、tf.TensorArray、tf.config分配GPU

TensorFlow 2.0 - TFRecord儲存資料集、@tf.function圖執行模式、tf.TensorArray、tf.config分配GPU

技術標籤:TensorFlow

文章目錄

學習於:簡單粗暴 TensorFlow 2

1. TFRecord 格式儲存

  • 使用該種格式,更高效地進行大規模的模型訓練

import random
import os
import tensorflow as tf

# 使用前一節 kaggle 上的 貓狗資料集
train_data_dir = "./dogs-vs-cats/train/"
test_data_dir =
"./dogs-vs-cats/test/" # 訓練檔案路徑 file_dir = [train_data_dir + filename for filename in os.listdir(train_data_dir)] labels = [0 if filename[0] == 'c' else 1 for filename in os.listdir(train_data_dir)] # 打包並打亂 f_l = list(zip(file_dir, labels)) random.shuffle(f_l) file_dir, labels = zip
(*f_l) # 切分訓練集,驗證集 valid_ratio = 0.1 idx = int((1 - valid_ratio) * len(file_dir)) train_files, valid_files = file_dir[:idx], file_dir[idx:] train_labels, valid_labels = labels[:idx], labels[idx:] # tfrecord 格式資料儲存路徑 train_tfrecord_file = "./dogs-vs-cats/train.tfrecords" valid_tfrecord_file =
"./dogs-vs-cats/valid.tfrecords" # -------------------看下面程式碼----------------------------- # 儲存過程 # 預先定義一個寫入器 with tf.io.TFRecordWriter(path=train_tfrecord_file) as writer: # 遍歷原始資料 for filename, label in zip(train_files, train_labels): img = open(filename, 'rb').read() # 讀取圖片,img 是 Byte 型別的字串 # 建立 feature 的 字典 k : v feature = { 'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img])), 'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])) } # feature 包裹成 example example = tf.train.Example(features=tf.train.Features(feature=feature)) # example 序列化為字串,寫入 writer.write(example.SerializeToString()) # -------------------看下面程式碼----------------------------- # 讀取過程 # 讀取 tfrecord 資料,得到 tf.data.Dataset 物件 raw_train_dataset = tf.data.TFRecordDataset(train_tfrecord_file) # 特徵的格式、資料型別 feature_description = { 'image': tf.io.FixedLenFeature(shape=[], dtype=tf.string), 'label': tf.io.FixedLenFeature([], tf.int64), } def _parse_example(example_string): # 解碼每個example # tf.io.parse_single_example 反序列化 feature_dict = tf.io.parse_single_example(example_string, feature_description) # 影象解碼 feature_dict['image'] = tf.io.decode_jpeg(feature_dict['image']) # 返回資料 X, y return feature_dict['image'], feature_dict['label'] # 處理資料集 train_dataset = raw_train_dataset.map(_parse_example) import matplotlib.pyplot as plt for img, label in train_dataset: plt.title('cat' if label==0 else 'dog') plt.imshow(img.numpy()) plt.show()

2. tf.function 高效能

  • TF 2.0 預設 即時執行模式(Eager Execution),靈活、易除錯
  • 追求高效能、部署模型時,使用圖執行模式(Graph Execution)
  • TF 2.0 的 tf.function 模組 + AutoGraph 機制,使用 @tf.function 修飾符,就可以將模型以圖執行模式執行

注意:@tf.function修飾的函式內,儘量只用 tf 的內建函式,變數只用 tensor、numpy 陣列

  • 被修飾的函式 F(X, y) 可以呼叫get_concrete_function 方法,獲得計算圖
graph = F.get_concrete_function(X, y)

3. tf.TensorArray 支援計算圖特性

  • tf.TensorArray 支援計算圖模式的 動態陣列
arr = tf.TensorArray(dtype=tf.int64, size=1, dynamic_size=True)
arr = arr.write(index=1, value=512)
# arr.write(index=0, value=512) # 沒有左值接受,會丟失
for i in range(arr.size()):
    print(arr.read(i))

4. tf.config 分配GPU

  • 列出裝置 list_physical_devices
print('---device----')
gpus = tf.config.list_physical_devices(device_type='GPU')
cpus = tf.config.list_physical_devices(device_type='CPU')
print(gpus, "\n", cpus)
# 單個的 GPU, CPU
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 
 [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]
  • 設定哪些可見 set_visible_devices
tf.config.set_visible_devices(devices=gpus[0:2], device_type='GPU')

或者

  • 終端輸入 export CUDA_VISIBLE_DEVICES=2,3
  • or 程式碼中加入
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"

指定程式 只在 顯示卡 2, 3 上執行

  • 視訊記憶體使用策略:
gpus = tf.config.list_physical_devices(device_type='GPU')
for gpu in gpus:
	# 僅在需要時申請視訊記憶體
    tf.config.experimental.set_memory_growth(device=gpu, enable=True)
gpus = tf.config.list_physical_devices(device_type='GPU')
# 固定視訊記憶體使用上限,超出報錯
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
  • 單 GPU 模擬多 GPU 環境

在單GPU電腦上,寫 多GPU 程式碼,可以模擬實現

gpus = tf.config.list_physical_devices('GPU')
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=2048),
     tf.config.LogicalDeviceConfiguration(memory_limit=2048)])
gpus = tf.config.list_logical_devices(device_type='GPU')
print(gpus)

輸出:2個虛擬的GPU

[LogicalDevice(name='/device:GPU:0', device_type='GPU'), 
 LogicalDevice(name='/device:GPU:1', device_type='GPU')]