问题遇到的现象和发生背景
用自己的数据集训练YOLOv4网络，安装好环境之后在cpu上运行正常，在gpu上报错：tensorflow.python.framework.errors_impl.InternalError: Blas SGEMM launch failed : m=173056, n=64, k=64 [Op:Conv2D]。我的环境：GPU：3080；驱动>460.6;CUDA:v11.2;cudnn:8.1.0;tensorflow-2.5.0(2.3.0,2.4.0,2.5.0都不行)
问题相关代码，请勿粘贴截图
import tensorflow as tf
from absl import app, flags, logging
from absl.flags import FLAGS
import os
import shutil
from core.yolov4 import YOLO, decode, compute_loss, decode_train
from core.dataset import Dataset
from core.config import cfg
import numpy as np
from core import utils
from core.utils import freeze_all, unfreeze_all
# gpus = tf.compat.v1.config.experimental.list_physical_devices('GPU')
# if gpus:
# try:
# for gpu in gpus:
# tf.compat.v1.config.experimental.set_memory_growth(gpu, True)
# logical_gpus = tf.compat.v1.config.experimental.list_logical_devices('GPU')
# except RuntimeError as e:
# print(e)
# print(tf.config.list_physical_devices('GPU'))
# print('Default GPU Device:{}'.format(tf.test.gpu_device_name()))
# print(tf.test.is_built_with_cuda())
# config = tf.compat.v1.ConfigProto()
# config.gpu_options.allow_growth = True
# config.gpu_options.per_process_gpu_memory_fraction = 0.5
# set_session(tf.compat.v1.Session(config=config))
flags.DEFINE_string('model', 'yolov4', 'yolov4, yolov3')
flags.DEFINE_string('weights', './scripts/yolov4.weights', 'pretrained weights')
flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny')
print('Parameters defined!')
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
# if len(physical_devices) > 0:
# tf.config.experimental.set_memory_growth(physical_devices[0], True)
trainset = Dataset(FLAGS, is_training=True)
# print(trainset)
testset = Dataset(FLAGS, is_training=False)
logdir = "./data/log"
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
# train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny)
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
if FLAGS.tiny:
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
else:
bbox_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
elif i == 1:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
else:
bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
if FLAGS.weights == None:
print("Training from scratch")
else:
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
else:
model.load_weights(FLAGS.weights)
print('Restoring weights from: %s ... ' % FLAGS.weights)
print('开始优化算法')
optimizer = tf.keras.optimizers.Adam()
print('优化完成')
if os.path.exists(logdir):
shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
# define training step function
# @tf.function
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d/%4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(),
giou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
(1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
)
optimizer.lr.assign(lr.numpy())
# writing summary data
with writer.as_default():
tf.summary.scalar("lr", optimizer.lr, step=global_steps)
tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
writer.flush()
def test_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(len(freeze_layers)):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
prob_loss, total_loss))
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in freeze_layers:
freeze = model.get_layer(name)
freeze_all(freeze)
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in freeze_layers:
freeze = model.get_layer(name)
unfreeze_all(freeze)
for image_data, target in trainset:
train_step(image_data, target)
for image_data, target in testset:
test_step(image_data, target)
model.save_weights("./checkpoints/yolov4")
if __name__ == '__main__':
try:
app.run(main)
except SystemExit:
pass
运行结果及报错内容
2022-04-01 20:40:02.771498: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cudnn64_8.dll
2022-04-01 20:40:03.675429: I tensorflow/stream_executor/cuda/cuda_dnn.cc:359] Loaded cuDNN version 8100
2022-04-01 20:40:05.048997: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cublas64_11.dll
2022-04-01 20:40:05.372659: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cublasLt64_11.dll
2022-04-01 20:40:06.180745: I tensorflow/stream_executor/cuda/cuda_blas.cc:1838] TensorFloat-32 will be used for the matrix multiplication. This will only be logged once.
2022-04-01 20:40:06.180912: E tensorflow/stream_executor/cuda/cuda_blas.cc:197] failed to set new cublas math mode: CUBLAS_STATUS_INVALID_VALUE
Traceback (most recent call last):
File "C:/Users/dell/Desktop/my_Program/tensorflow-yolov4-tflite-master/train.py", line 190, in
app.run(main)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\absl\app.py", line 312, in run
_run_main(main, args)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\absl\app.py", line 258, in _run_main
sys.exit(main(argv))
File "C:/Users/dell/Desktop/my_Program/tensorflow-yolov4-tflite-master/train.py", line 182, in main
train_step(image_data, target)
File "C:/Users/dell/Desktop/my_Program/tensorflow-yolov4-tflite-master/train.py", line 109, in train_step
pred_result = model(image_data, training=True)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1030, in call
outputs = call_fn(inputs, *args, **kwargs)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\keras\engine\functional.py", line 421, in call
inputs, training=training, mask=mask)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\keras\engine\functional.py", line 556, in _run_internal_graph
outputs = node.layer(*args, **kwargs)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1030, in call
outputs = call_fn(inputs, *args, **kwargs)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\keras\layers\convolutional.py", line 249, in call
outputs = self._convolution_op(inputs, self.kernel)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\util\dispatch.py", line 206, in wrapper
return target(*args, **kwargs)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\ops\nn_ops.py", line 1019, in convolution_v2
name=name)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\ops\nn_ops.py", line 1149, in convolution_internal
name=name)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\ops\nn_ops.py", line 2603, in _conv2d_expanded_batch
name=name)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\ops\gen_nn_ops.py", line 931, in conv2d
_ops.raise_from_not_ok_status(e, name)
File "C:\Users\dell\Desktop\my_Program\tensorflow-yolov4-tflite-master\venv\lib\site-packages\tensorflow\python\framework\ops.py", line 6897, in raise_from_not_ok_status
six.raise_from(core._status_to_exception(e.code, message), None)
File "", line 3, in raise_from
tensorflow.python.framework.errors_impl.InternalError: Blas SGEMM launch failed : m=173056, n=64, k=64 [Op:Conv2D]
我的解答思路和尝试过的方法
已经试过设置代码占用一定比例GPU或者动态申请GPU，然而没用。
我想要达到的结果：烦请各路同志不吝赐教！
