Tensorflow ImageAI GPU使用情况

时间:2020-06-03 20:05:43

标签: python tensorflow gpu imageai

我使用imageai裁剪我的汽车数据集以进行cnn训练。我的代码很简单,我有两个GPU,并编码到张量流日志和nvidia-smi输出,它使用了我的两个GPU,但同时没有使用所有可用的内存

import os
import sys, getopt
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
from imageai.Detection import ObjectDetection
from time import perf_counter
import shutil
import numpy as np
from io import BytesIO
from PIL import Image
import argparse
import time
from shutil import copyfile
parser = argparse.ArgumentParser(description='Car Images Filter')

parser.add_argument("-i", "--input_folder", action="store", dest="input_folder",   help="Folder for filtering", type=str, required=True)
parser.add_argument("-o", "--output_folder", action="store", dest="output_folder",   help="Folder for the results of wrong images", type=str, required=False)
parser.add_argument("-start_from_model", "--start_from_model", action="store", dest="start_from_model",   help="Start model for exception cases", type=str, default='')
parser.add_argument("-start_from_mark", "--start_from_mark", action="store", dest="start_from_mark",   help="Start mark for exception cases", type=str, default='')
parser.add_argument("-start_from_file", "--start_from_file", action="store", dest="start_from_file",   help="Start file for exception cases", type=str, default='')
parser.add_argument("-remove", "--remove", action="store_true", dest="remove",   help="Folder for the results of wrong images", required=False)


args = parser.parse_args()
execution_path = os.getcwd()
detector = ObjectDetection()
detector.setModelTypeAsRetinaNet()
detector.setModelPath(os.path.join(execution_path , "resnet50_coco_best_v2.0.1.h5"))
detector.loadModel()
directories = os.listdir(args.input_folder)
if args.start_from_mark:
    directories = directories[directories.index(args.start_from_mark):]
for mark in directories:
    models = os.listdir(os.path.join(args.input_folder, mark))
    for model in models:
        files = os.listdir(os.path.join(args.input_folder, mark, model))
        if model==args.start_from_model and args.start_from_file:
            files=files[files.index(args.start_from_file):]
        for file in files:
            if(file.endswith('.json')):
                continue
            image = Image.open(os.path.join(args.input_folder, mark, model, file))
            detected_image_array, detections = detector.detectCustomObjectsFromImage(custom_objects=detector.CustomObjects(car=True,truck=True), input_image=os.path.join(args.input_folder, mark, model, file), output_type="array")
            detected_image = Image.fromarray(detected_image_array, 'RGB')
            if len(detections) == 0:
                copyfile(os.path.join(args.input_folder, mark, model, file), os.path.join(args.output_folder, mark, model, file))
            print(mark + " " + model + " " + file)

这是我的nvidia-smi输出: Image link 如何增加GPU内存使用率以使其更快?

0 个答案:

没有答案