欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

防止keras占满GPU显存-keras-yolo

程序员文章站 2022-07-12 12:32:11
...
import os
import numpy as np
import keras.backend as K
from keras.layers import Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
# 指定那几块GPU是对于本程序可见的 其实有没有都没有什么区别 tensor默认占用全部的
# 一般用在忽略哪块GPU 多GPU只能用一块那种任务
from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
from yolo3.utils import get_random_data
# 多Gpu训练
from keras.utils import multi_gpu_model
import keras.backend.tensorflow_backend as KTF
import tensorflow as tf

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

def _main():
    # 训练前记得修改yolov3.cfg和model_data/voc_class.txt
    annotation_path = '1.14/train/label1.txt' # 这里是要的由VOC生成的txt,里面包括image地址和框和种类 需修改
    log_dir = '1.14/weights/' # 是存放生成好的新的weights的
    classes_path = '1.14/train/classes.txt' # 这个我只标记了4类,所以要修改,原来voc是有20类的
    anchors_path = '1.14/train/anchors.txt' # 这里的anchor暂时先不修改(最好从新kmean算anchors)
    class_names = get_classes(classes_path)
    num_classes = len(class_names)
    anchors = get_anchors(anchors_path)

    # 修改大小,按照图片最小的边来修改
    input_shape = (544,960) # multiple of 32, hw 这里是否需要修改

    config = tf.ConfigProto()
    config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
    sess = tf.Session(config=config)

    KTF.set_session(sess)

    is_tiny_version = len(anchors)==6 # default setting
    if is_tiny_version:
        model = create_tiny_model(input_shape, anchors, num_classes,
            freeze_body=2, weights_path='model_data/tiny_yolo_weights.h5')
    else:
        # 这里要修改,这里的权重是基础权重,在这个基础上继续训练
        model = create_model(input_shape, anchors, num_classes,
            freeze_body=2, weights_path='12.5/weigths/dog_detect_544x960.h5') # make sure you know what you freeze

 

相关标签: 无人驾驶