文章目錄

• • 1. TFRecord 格式存儲
  • 2. tf.function 高性能
  • 3. tf.TensorArray 支持計算圖特性
  • 4. tf.config 分配GPU

學習于：簡單粗暴 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 1for 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 : vfeature = {'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img])),'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))}# feature 包裹成 exampleexample = 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, yreturn 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')]

總結

以上是生活随笔為你收集整理的TensorFlow 2.0 - TFRecord存储数据集、@tf.function图执行模式、tf.TensorArray、tf.config分配GPU的全部內容，希望文章能夠幫你解決所遇到的問題。

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