keras 使用单个或多个GPU训练模型（已经不需要自己设置了，会自动检测环境）

最新推荐文章于 2023-04-26 10:12:08 发布

云月月白

最新推荐文章于 2023-04-26 10:12:08 发布

阅读量1.1k

点赞数 5

分类专栏： Python keras GPU 文章标签： gpu

本文链接：https://blog.csdn.net/Le___Le/article/details/110672347

版权

本篇博客介绍如何在Keras中利用单个或多个GPU进行模型训练，无需手动配置，系统将自动检测GPU环境。对于单GPU训练，环境设置简单；而对于多GPU训练，则需额外指定GPU数量。

摘要由CSDN通过智能技术生成

问题： keras 使用单个或多个GPU训练模型

解决办法：

单个GPU：直接设置环境

import os
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 指定GPU编号

def get_model():
	# input_ids = keras.layers.Input(shape=(MAXLEN, ))...
	#_, x = bert_model([input_ids, attension_mask, token_type_ids])
	# outputs = keras.layers.Dense(1, activation='sigmoid')(x)
	
    model = keras.models.Model(inputs=inputs, outputs=outputs)
    model.compile(loss='binary_crossentropy', optimizer=keras.optimizers.Adam(lr=LR), metrics=['accuracy'])
    return model

多个GPU：设置环境 + 规定数目

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0， 1， 2， 3" # 改变编号

def get_model():
	# input_ids = keras.layers.Input(shape=(MAXLEN, ))...
	#_, x = bert_model([input_ids, attension_m