python - 使用 TensorFlow 在 google colab 上同时使用 CPU 和 GPU

Question

我有一个文件，里面有 1000 万个 float 。使用合并排序(下面的客户代码)，我想对这个文件进行排序。在这里，我一次读取 100 万条数据并将数字排序并写入一个单独的文件以便稍后合并。

使用 TensorFlow 2.0 API 自定义以下 Mergesort 代码。

import tensorflow as tf
#@tf.function
def split_list(input_list):
    input_list_len = len(input_list)
    midpoint = input_list_len // 2
    return input_list[:midpoint], input_list[midpoint:]

#@tf.function
def merge_sorted_lists(list_left, list_right):

    if tf.math.equal(len(list_left), 0) :
        return list_right
    elif tf.math.equal(len(list_right), 0):
        return list_left
    
 
    index_left = index_right = 0
    list_merged = []  # list to build and return
    list_len_target = len(list_left) + len(list_right)
    while tf.math.less(len(list_merged), list_len_target):
        if tf.math.less_equal(list_left[index_left], list_right[index_right]) :
            # Value on the left list is smaller (or equal so it should be selected)
            list_merged = [*list_merged,list_left[index_left]]
            index_left += 1
        else:
            # Right value bigger
            list_merged = [*list_merged,list_right[index_right]]
            index_right += 1
            

        if tf.math.equal(index_right, len(list_right)):
            list_merged = [*list_merged,*list_left[index_left:]]
            break
        elif tf.math.equal(index_left, len(list_left)):
            list_merged = [*list_merged,*list_right[index_right:]]
            break
    return list_merged

#@tf.function
def merge_sort(input_list):
    if tf.math.less_equal(len(input_list), 1):
        return input_list
    else:
        left, right = split_list(input_list)
        return merge_sorted_lists(merge_sort(left), merge_sort(right))

下面的代码在所需的主机/设备上调用上面的 merge_sort 函数。

def cpu(input_list):
  lines = []
  with tf.device('/cpu:0'):
    lines = merge_sort(input_list)
    lines1 = []
    lines1 = [str(x)+"\n" for x in lines]
  return lines1

def GPU(input_list):
  lines = []
  with tf.device('/device:GPU:0'):
    lines = merge_sort(input_list)
    lines1 = []
    lines1 = [str(x)+"\n" for x in lines]
  return lines1

上述 CPU 或 GPU 函数调用如下，使用迭代器一次读取 100 万:

dataset = tf.data.TextLineDataset("numbers.txt") #10 million rows in single file
dataset = dataset.batch(1_000_000)  # divide into 10 batches of 1million each
dataset = dataset.map(lambda x: tf.strings.to_number(tf.strings.strip(x), tf.float32))
iterator = dataset.__iter__()
    
start_time = timeit.default_timer()
lines = cpu(np.stack(list(iterator.get_next())))
print(lines)
print("cpu_time ", timeit.default_timer() - start_time)

start_time = timeit.default_timer()
lines = GPU(np.stack(list(iterator.get_next())))
print(lines)
print("gpu_time ", timeit.default_timer() - start_time)

fid = 1
f_out = open('chunk_{}.txt'.format(fid), 'w')
f_out.writelines(lines)

以上代码运行Google Colab GPU 实例，前 100 万条记录在 CPU 上排序一次，第二百万条记录在 GPU 上排序。但是我想同时使用 CPU 和 GPU 对不同的百万数字集运行合并排序。暂时忽略合并。我将不胜感激，示例代码行(如果存在，几行)或指南(如果涉及大量工作)，以实现这一目标。

Answer 1

您可以使用内置分布策略在多个 GPU、多台机器或 TPU 之间分布训练。

MultiWorkerMirroredStrategy 就是您正在寻找的那个。它有两种跨设备通信的实现。 CommunicationImplementation.RING 基于 RPC，同时支持 CPU 和 GPU。

https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/CommunicationImplementation#RING