使用 TensorFlow 在 google colab 上同时使用 CPU 和 GPU答案

【问题标题】：use both CPU and GPU on google colab parallelly, using TensofFlow使用 TensorFlow 在 google colab 上同时使用 CPU 和 GPU
【发布时间】：2021-03-10 07:56:36
【问题描述】：

我有一个文件，里面有 1000 万个浮点数。使用合并排序（下面的自定义代码），我想对这个文件进行排序。在这里，我一次读取 100 万个并排序数字并写入一个单独的文件以便稍后合并。

使用 TensorFlow 2.0 API 自定义以下合并排序代码。

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 在不同的百万数字集上运行归并排序。暂时忽略合并。我将不胜感激，示例代码行（如果存在，几行）或指南（如果涉及大量工作），以实现这一目标。

【问题讨论】：

使用 Tensorflow distribution strategy 库在不同的运行时（GPU、CPU）上并行运行您的代码。谢谢！

标签： python tensorflow gpu google-colaboratory mergesort

【解决方案1】：

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

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

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

【讨论】：