我使用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内存使用率以使其更快?