Linux -- 使用多张gpu卡进行深度学习任务(以tensorflow为例)

在linux系统上进行多gpu卡的深度学习任务

  • 确保已安装最新的 TensorFlow GPU 版本。
python 复制代码
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
  • 1、确保你已经正确安装了tensorflow和相关的GPU驱动,这里可以通过在命令行输入nvidia-smi来查看:

    如果成功显示了类似上述的GPU信息和驱动版本信息,则说明NVIDIA驱动已经正确安装。

2、导入必要的库,设置可见的gpu设备列表:

python 复制代码
import tensorflow as tf
# 设置可见的GPU设备列表(例如,使用GPU 0、1、2和3)
gpu_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(gpu_devices, 'GPU')
  • 3、创建一个MirroredStrategy对象,该对象将自动复制模型和数据到每个可见的GPU卡上:
python 复制代码
strategy = tf.distribute.MirroredStrategy()
  • 4、在strategy范围内创建和训练模型:
python 复制代码
with strategy.scope():
    # 创建和编译模型
    model = create_model()
    model.compile(...)
    
    # 加载数据
    train_dataset = load_train_data()
    test_dataset = load_test_data()
    
    # 训练模型
    model.fit(train_dataset, validation_data=test_dataset, ...)

以上,在MirroredStrategy范围内创建的模型将自动复制并分布到每个可见的GPU卡上,每个卡都将处理一部分数据。

使用多个 GPU 的最佳做法是使用 tf.distribute.Strategy

以下给出一个官网的简单示例:

python 复制代码
tf.debugging.set_log_device_placement(True)
gpus = tf.config.list_logical_devices('GPU')
strategy = tf.distribute.MirroredStrategy(gpus)
with strategy.scope():
  inputs = tf.keras.layers.Input(shape=(1,))
  predictions = tf.keras.layers.Dense(1)(inputs)
  model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
  model.compile(loss='mse',
                optimizer=tf.keras.optimizers.SGD(learning_rate=0.2))

当然,也有手动的放置方法:

python 复制代码
tf.debugging.set_log_device_placement(True)

gpus = tf.config.list_logical_devices('GPU')
if gpus:
  # Replicate your computation on multiple GPUs
  c = []
  for gpu in gpus:
    with tf.device(gpu.name):
      a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
      c.append(tf.matmul(a, b))

  with tf.device('/CPU:0'):
    matmul_sum = tf.add_n(c)

  print(matmul_sum)

在tensorflow上使用gpu:https://www.tensorflow.org/guide/gpu?hl=zh-cn

相关推荐
Yan-英杰4 分钟前
百度搜索和文心智能体接入DeepSeek满血版——AI搜索的新纪元
图像处理·人工智能·python·深度学习·deepseek
程序员yt4 分钟前
双非一本电子信息专业自学嵌入式,学完 Linux 后咋走?单片机 & FreeRTOS 要补吗?
linux·运维·单片机
taoqick2 小时前
对PosWiseFFN的改进: MoE、PKM、UltraMem
人工智能·pytorch·深度学习
安於宿命5 小时前
【Linux】进程间通信——进程池
linux·c++
新兴ICT项目支撑5 小时前
天翼云910B部署DeepSeek蒸馏70B LLaMA模型实践总结
linux·运维·服务器·910b·天翼云·deepseek r1
charles_vaez5 小时前
开源模型应用落地-LangGraph101-探索 LangGraph 短期记忆
深度学习·语言模型·自然语言处理
WHATEVER_LEO6 小时前
【每日论文】Latent Radiance Fields with 3D-aware 2D Representations
人工智能·深度学习·神经网络·机器学习·计算机视觉·自然语言处理
Kai HVZ7 小时前
《深度学习》——调整学习率和保存使用最优模型
人工智能·深度学习·学习
Spike()7 小时前
HaProxy源码安装(Rocky8)
linux·负载均衡
网硕互联的小客服7 小时前
如何提高网站在百度中的权重?
linux·运维·服务器·windows·安全