first commit
This commit is contained in:
Semi212 committed
commit
54ccfb11fe
218 files changed
+19525
No files matched your search
Executable
Whitespace-only changes.
Executable
+422
@@ -0,0 +1,422 @@
|
||||
import logging
|
||||
import os
|
||||
from vizer.draw import draw_boxes
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from tqdm import tqdm
|
||||
from ssd.config import cfg
|
||||
from ssd.data.build import make_data_loader
|
||||
from ssd.data.datasets.evaluation import evaluate
|
||||
import datetime
|
||||
from ssd.utils import dist_util, mkdir
|
||||
from ssd.utils.dist_util import synchronize, is_main_process
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
import logging
|
||||
import time
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from ssd.utils.dist_util import is_main_process, get_world_size
|
||||
from ssd.utils.dist_util import synchronize
|
||||
from ssd.utils.timer import Timer, get_time_str
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu):
|
||||
all_predictions = dist_util.all_gather(predictions_per_gpu)
|
||||
if not dist_util.is_main_process():
|
||||
return
|
||||
# merge the list of dicts
|
||||
predictions = {}
|
||||
for p in all_predictions:
|
||||
predictions.update(p)
|
||||
# convert a dict where the key is the index in a list
|
||||
image_ids = list(sorted(predictions.keys()))
|
||||
if len(image_ids) != image_ids[-1] + 1:
|
||||
logger = logging.getLogger("SSD.inference")
|
||||
logger.warning(
|
||||
"Number of images that were gathered from multiple processes is not "
|
||||
"a contiguous set. Some images might be missing from the evaluation"
|
||||
)
|
||||
|
||||
# convert to a list
|
||||
predictions = [predictions[i] for i in image_ids]
|
||||
return predictions
|
||||
|
||||
|
||||
def compute_on_dataset_byHatcherW(model, data_loader, device,timer=None):
|
||||
results_dict = {}
|
||||
cpu_device = torch.device("cpu")
|
||||
now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||||
filename = os.path.join(
|
||||
"/home/baai/voc/HatcherW/SSD/", now_time)
|
||||
isExists = os.path.exists(filename)
|
||||
if not isExists:
|
||||
os.makedirs(filename)
|
||||
print(filename + ' 创建成功')
|
||||
else:
|
||||
print(filename + ' 目录已存在')
|
||||
with open(
|
||||
"/home/baai/voc/HatcherW/ImageSets/Main/test.txt",
|
||||
"r") as f:
|
||||
for batch in tqdm(data_loader):
|
||||
# 每次进循环先读取一次文件路径
|
||||
|
||||
|
||||
JPEG_NAME = f.readline().strip('\n')
|
||||
JPEG_NAME = JPEG_NAME + '.jpg'
|
||||
print(JPEG_NAME)
|
||||
filename = os.path.join(
|
||||
"/home/baai/voc/HatcherW/JPEGImages/",
|
||||
JPEG_NAME)
|
||||
print(filename)
|
||||
####################################
|
||||
# src = cv2.imread(filename,flags=1)
|
||||
# cv2.namedWindow("yuanlai", cv2.WINDOW_NORMAL)
|
||||
# cv2.imshow("yuanlai", src)
|
||||
# contrast_brightness_image(src, 0.5, 5) # 第一个1.2为对比度 第二个为亮度数值越大越亮
|
||||
|
||||
|
||||
|
||||
########################################
|
||||
|
||||
|
||||
image_forinf = Image.open(filename).convert("RGB") #read image
|
||||
image_forinf = np.array(image_forinf)[:, :, [2, 1, 0]] #conver image into array
|
||||
# cv2.imshow('测试', image_forinf)
|
||||
images, target, image_ids = batch # target in inference does not be used
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
if timer:
|
||||
timer.tic()
|
||||
output = model(images.to(device))
|
||||
if timer:
|
||||
if not cfg.MODEL.DEVICE == 'cpu':
|
||||
torch.cuda.synchronize()
|
||||
timer.toc()
|
||||
|
||||
output = [o.to(cpu_device) for o in output]
|
||||
##########inference 测试,用于测试如何提取数据
|
||||
print("output is {}",format(output))
|
||||
output_temp = output[0]
|
||||
# print("output 类型 is {}".format(type(output[0]['scores'][0])))
|
||||
# print("output fenshu is {}".format((output[0]['scores'][0])))
|
||||
# print("output biankuang is {}".format((output_temp['boxes'][0,:])))
|
||||
# print("output biankuangnum is {}".format((output_temp['boxes'].size())))
|
||||
#print("output label is {}".format((output_temp['labels'][0])))
|
||||
# print("hhhh")
|
||||
bbox_nums = output_temp['boxes'].size()[0] #取出box的个数
|
||||
output_scores = output[0]['scores'].numpy()
|
||||
output_labels = output[0]['labels'].numpy()
|
||||
output_boxes = output[0]['boxes'].numpy()
|
||||
|
||||
print("output 类型 is {}".format(type(output_scores)))
|
||||
result = image_forinf.copy()
|
||||
result1 = image_forinf.copy()
|
||||
heigh = image_forinf.shape[0]
|
||||
kuan = image_forinf.shape[1]
|
||||
print(heigh,kuan)
|
||||
###visual by hcwang
|
||||
for bbox_num in range(bbox_nums): # 画框 需要画bbox_nums个
|
||||
left_top = (
|
||||
(int)(output_boxes[bbox_num][0] * kuan / 300.0), (int)(output_boxes[bbox_num][1] * heigh / 300.0))
|
||||
right_bottom = (
|
||||
(int)(output_boxes[bbox_num][2] * kuan / 300.0), (int)(output_boxes[bbox_num][3] * heigh / 300.0))
|
||||
cv2.rectangle(result, left_top, right_bottom, (0, 0, 255), 3)
|
||||
#cv2.imshow('测试',image_forinf)
|
||||
# 存储
|
||||
filename = os.path.join("/home/baai/voc/HatcherW/SSD/", now_time ,JPEG_NAME)
|
||||
# cv2.imwrite(filename, result)
|
||||
for bbox_num in range(bbox_nums):
|
||||
output_boxes[bbox_num][0] *=kuan / 300.0
|
||||
output_boxes[bbox_num][1] *=heigh / 300.0
|
||||
output_boxes[bbox_num][2] *=kuan / 300.0
|
||||
output_boxes[bbox_num][3] *=heigh / 300.0
|
||||
class_names = ('__background__',
|
||||
'aeroplane', 'bicycle', 'bird', 'boat',
|
||||
'bottle', 'bus', 'car', 'cat', 'chair',
|
||||
'cow', 'diningtable', 'dog', 'horse',
|
||||
'motorbike', 'person', 'pottedplant',
|
||||
'sheep', 'sofa', 'train', 'tvmonitor')
|
||||
drawn_image = draw_boxes(result1, output_boxes, output_labels, output_scores, class_names).astype(np.uint8)
|
||||
Image.fromarray(drawn_image).save(filename)
|
||||
# cv2.imshow('result',drawn_image) #显示结果
|
||||
# cv2.waitKey(0)
|
||||
# cv2.destroyAllWindows()
|
||||
|
||||
print(output)
|
||||
results_dict.update(
|
||||
{img_id: result for img_id, result in zip(image_ids, output)}
|
||||
)
|
||||
f.close()
|
||||
return results_dict
|
||||
def compute_on_dataset(model, data_loader, device):
|
||||
results_dict = {}
|
||||
for batch in tqdm(data_loader):
|
||||
images, targets, image_ids = batch
|
||||
cpu_device = torch.device("cpu")
|
||||
with torch.no_grad():
|
||||
outputs = model(images.to(device))
|
||||
|
||||
outputs = [o.to(cpu_device) for o in outputs]
|
||||
results_dict.update(
|
||||
{img_id: result for img_id, result in zip(image_ids, outputs)}
|
||||
)
|
||||
return results_dict
|
||||
|
||||
def inference_byHatcherW(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
|
||||
dataset = data_loader.dataset
|
||||
logger = logging.getLogger("SSD.inference")
|
||||
logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
|
||||
total_timer = Timer()
|
||||
inference_timer = Timer()
|
||||
total_timer.tic()
|
||||
|
||||
predictions_path = os.path.join(output_folder, 'predictions.pth')
|
||||
if use_cached and os.path.exists(predictions_path):
|
||||
predictions = torch.load(predictions_path, map_location='cpu')
|
||||
else:
|
||||
predictions = compute_on_dataset_byHatcherW(model, data_loader, device,inference_timer)
|
||||
synchronize()
|
||||
predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
||||
if not is_main_process():
|
||||
return
|
||||
if output_folder:
|
||||
torch.save(predictions, predictions_path)
|
||||
#return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
|
||||
# def inference_byHatcherW(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
|
||||
# dataset = data_loader.dataset
|
||||
# logger = logging.getLogger("SSD.inference")
|
||||
# logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
|
||||
# total_timer = Timer()
|
||||
# inference_timer = Timer()
|
||||
# total_timer.tic()
|
||||
# predictions_path = os.path.join(output_folder, 'predictions.pth')
|
||||
# if use_cached and os.path.exists(predictions_path):
|
||||
# predictions = torch.load(predictions_path, map_location='cpu')
|
||||
# else:
|
||||
# predictions = compute_on_dataset(model, data_loader, device,inference_timer)
|
||||
# synchronize()
|
||||
# predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
||||
# if not is_main_process():
|
||||
# return
|
||||
# if output_folder:
|
||||
# torch.save(predictions, predictions_path)
|
||||
# return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
|
||||
|
||||
def inference(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
|
||||
dataset = data_loader.dataset
|
||||
logger = logging.getLogger("SSD.inference")
|
||||
logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
|
||||
predictions_path = os.path.join(output_folder, 'predictions.pth')
|
||||
if use_cached and os.path.exists(predictions_path):
|
||||
predictions = torch.load(predictions_path, map_location='cpu')
|
||||
else:
|
||||
predictions = compute_on_dataset(model, data_loader, device)
|
||||
synchronize()
|
||||
predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
||||
if not is_main_process():
|
||||
return
|
||||
if output_folder:
|
||||
torch.save(predictions, predictions_path)
|
||||
return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
|
||||
@torch.no_grad()
|
||||
def do_evaluation_byHatcherW(cfg, model, distributed, **kwargs):
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
model = model.module
|
||||
model.eval()
|
||||
device = torch.device(cfg.MODEL.DEVICE)
|
||||
data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed)
|
||||
eval_results = []
|
||||
for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val):
|
||||
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
|
||||
if not os.path.exists(output_folder):
|
||||
mkdir(output_folder)
|
||||
eval_result = inference_byHatcherW(model, data_loader, dataset_name, device, output_folder, **kwargs)
|
||||
eval_results.append(eval_result)
|
||||
|
||||
|
||||
def do_evaluation(cfg, model, distributed, **kwargs):
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
model = model.module
|
||||
model.eval()
|
||||
device = torch.device(cfg.MODEL.DEVICE)
|
||||
data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed)
|
||||
eval_results = []
|
||||
for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val):
|
||||
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
|
||||
if not os.path.exists(output_folder):
|
||||
mkdir(output_folder)
|
||||
eval_result = inference(model, data_loader, dataset_name, device, output_folder, **kwargs)
|
||||
eval_results.append(eval_result)
|
||||
return eval_results
|
||||
# def inference_hcw(
|
||||
# model,
|
||||
# data_loader,
|
||||
# dataset_name,
|
||||
# iou_types=("bbox",),
|
||||
# box_only=False,
|
||||
# device="cuda",
|
||||
# expected_results=(),
|
||||
# expected_results_sigma_tol=4,
|
||||
# output_folder=None,
|
||||
# ):
|
||||
# # convert to a torch.device for efficiency
|
||||
# device = torch.device(device) #use GPU
|
||||
# num_devices = get_world_size() #get the num of GPU
|
||||
# logger = logging.getLogger("maskrcnn_benchmark.inference")
|
||||
# dataset = data_loader.dataset #get dataset
|
||||
# logger.info("Start evaluation on {} dataset({} images).".format(dataset_name, len(dataset)))
|
||||
# total_timer = Timer()
|
||||
# inference_timer = Timer()
|
||||
# total_timer.tic()
|
||||
# predictions = compute_on_dataset(model, data_loader, device, inference_timer)
|
||||
# # wait for all processes to complete before measuring the time
|
||||
# synchronize()
|
||||
# total_time = total_timer.toc()
|
||||
# total_time_str = get_time_str(total_time)
|
||||
# logger.info(
|
||||
# "Total run time: {} ({} s / img per device, on {} devices)".format(
|
||||
# total_time_str, total_time * num_devices / len(dataset), num_devices
|
||||
# )
|
||||
# )
|
||||
# total_infer_time = get_time_str(inference_timer.total_time)
|
||||
# logger.info(
|
||||
# "Model inference time: {} ({} s / img per device, on {} devices)".format(
|
||||
# total_infer_time,
|
||||
# inference_timer.total_time * num_devices / len(dataset),
|
||||
# num_devices,
|
||||
# )
|
||||
# )
|
||||
#
|
||||
# predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
||||
# #visual_whc(predictions)
|
||||
# if not is_main_process():
|
||||
# return
|
||||
#
|
||||
# if output_folder:
|
||||
# torch.save(predictions, os.path.join(output_folder, "predictions.txt"))
|
||||
#
|
||||
# extra_args = dict(
|
||||
# box_only=box_only,
|
||||
# iou_types=iou_types,
|
||||
# expected_results=expected_results,
|
||||
# expected_results_sigma_tol=expected_results_sigma_tol,
|
||||
# )
|
||||
# def compute_on_dataset_hcw(model, data_loader, device, timer=None):
|
||||
# model.eval() #不启用BN和DP
|
||||
# results_dict = {} #the result
|
||||
# cpu_device = torch.device("cpu")
|
||||
# with open(
|
||||
# "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/ImageSets/Main/test.txt",
|
||||
# "r") as f:
|
||||
# for _, batch in enumerate(tqdm(data_loader)):
|
||||
# #每次进循环先读取一次文件路径
|
||||
# JPEG_NAME = f.readline().strip('\n')
|
||||
# JPEG_NAME = JPEG_NAME + '.jpg'
|
||||
# print(JPEG_NAME)
|
||||
# filename = os.path.join(
|
||||
# "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/JPEGImages/",
|
||||
# JPEG_NAME)
|
||||
# print(filename)
|
||||
# image_forinf = Image.open(filename).convert("RGB")
|
||||
# image_forinf = np.array(image_forinf)[:, :, [2, 1, 0]]
|
||||
# # print("batch is "+batch)
|
||||
# images, image_ids = batch
|
||||
# # print("image_id is "+image_ids)
|
||||
# with torch.no_grad(): #确定是否反向传播求dao
|
||||
# if timer:
|
||||
# timer.tic()
|
||||
# output = model(images.to(device))
|
||||
# if timer:
|
||||
# if not cfg.MODEL.DEVICE == 'cpu':
|
||||
# torch.cuda.synchronize()
|
||||
# timer.toc()
|
||||
# output = [o.to(cpu_device) for o in output]
|
||||
# output_forinf = output[0]
|
||||
# print("output caculate success")
|
||||
# #将结果rezize
|
||||
# height, width = image_forinf.shape[:-1]
|
||||
# output_forinf = output_forinf.resize((width, height))
|
||||
# print("outout resize success")
|
||||
#
|
||||
# ################
|
||||
# top_predictions = select_top_predictions(output_forinf)
|
||||
# print("nms success")
|
||||
# result = image_forinf.copy()
|
||||
# result = overlay_boxes(result, top_predictions)
|
||||
# result = overlay_class_names(result, top_predictions)
|
||||
# print("print reslut handle success")
|
||||
# ################
|
||||
# filename = os.path.join(
|
||||
# "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/result1/", JPEG_NAME)
|
||||
# cv2.imwrite(filename, result)
|
||||
# results_dict.update(
|
||||
# {img_id: result for img_id, result in zip(image_ids, output)}
|
||||
# )
|
||||
# ##为了测试num_bboxex
|
||||
# # print("output1:")
|
||||
# # temp = results_dict[0].bbox.size()
|
||||
# # print(temp[0])
|
||||
# # print(temp[1])
|
||||
# # print(results_dict[0].bbox.size())
|
||||
# #
|
||||
# #
|
||||
# # print(type(results_dict[0].bbox))
|
||||
#
|
||||
# # with open(
|
||||
# # "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/ImageSets/Main/test.txt",
|
||||
# # "r") as f:
|
||||
# # JPEG_NAME = f.readline().strip('\n')
|
||||
# # # JPEG_NAME.strip('\n')
|
||||
# # JPEG_NAME = JPEG_NAME + '.jpg'
|
||||
# # print(JPEG_NAME)
|
||||
# # f.close()
|
||||
# # filename = os.path.join("/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/JPEGImages/", JPEG_NAME)
|
||||
# # print(filename)
|
||||
# # img = cv2.imread(filename)
|
||||
# # #cv2.imshow('测试', img)
|
||||
# # #cv2.waitKey(0)
|
||||
# # #cv2.destroyAllWindows() # dv2.destroyWindow(wname)
|
||||
# # results_dict_np = results_dict[0].bbox.numpy()
|
||||
# # left_top = ((int)(results_dict_np[0][0]*1920.0/1333.0), (int)(results_dict_np[0][1]*1920.0/1333.0))
|
||||
# # right_bottom = ((int)(results_dict_np[0][2]*1920.0/1333.0), (int)(results_dict_np[0][3]*1920.0/1333.0))
|
||||
# # cv2.rectangle(img, left_top, right_bottom, (0, 0, 255), 3)
|
||||
# # print(results_dict_np[2][1])
|
||||
# # filename = os.path.join("/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/result/", JPEG_NAME)
|
||||
# # cv2.imwrite(filename, img)
|
||||
# print('hhhh')
|
||||
# print('hhhh')
|
||||
# #visual_whc(results_dict)
|
||||
# f.close()
|
||||
# return results_dict
|
||||
def do_evaluation_hcw(cfg, model, distributed, **kwargs):
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
model = model.module
|
||||
model.eval()
|
||||
device = torch.device(cfg.MODEL.DEVICE)
|
||||
data_loaders_vals = make_data_loader(cfg, is_train=False, distributed=distributed)
|
||||
eval_results = []
|
||||
output_folders = [None] * len(cfg.DATASETS.TEST)
|
||||
if cfg.OUTPUT_DIR:
|
||||
for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
|
||||
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
|
||||
mkdir(output_folder)
|
||||
output_folders[idx] = output_folder
|
||||
for output_folder, data_loaders_val in zip(output_folders,cfg.DATASETS.TEST, data_loaders_vals):
|
||||
eval_result = inference(model, data_loaders_val, dataset_name, device, output_folder, **kwargs)
|
||||
eval_results.append(eval_result)
|
||||
return eval_results
|
||||
|
||||
|
||||
def contrast_brightness_image(src1, a, g):
|
||||
h, w, ch = src1.shape # 获取shape的数值，height和width、通道
|
||||
|
||||
# 新建全零图片数组src2,将height和width，类型设置为原图片的通道类型(色素全为零，输出为全黑图片)
|
||||
src2 = np.zeros([h, w, ch], src1.dtype)
|
||||
dst = cv2.addWeighted(src1, a, src2, 1 - a, g) # addWeighted函数说明如下
|
||||
cv2.imshow("con-bri-demo", dst)
|
||||
|
||||
Executable
+140
@@ -0,0 +1,140 @@
|
||||
import collections
|
||||
import datetime
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ssd.engine.inference import do_evaluation
|
||||
from ssd.utils import dist_util
|
||||
from ssd.utils.metric_logger import MetricLogger
|
||||
|
||||
|
||||
def write_metric(eval_result, prefix, summary_writer, global_step):
|
||||
for key in eval_result:
|
||||
value = eval_result[key]
|
||||
tag = '{}/{}'.format(prefix, key)
|
||||
if isinstance(value, collections.Mapping):
|
||||
write_metric(value, tag, summary_writer, global_step)
|
||||
else:
|
||||
summary_writer.add_scalar(tag, value, global_step=global_step)
|
||||
|
||||
|
||||
def reduce_loss_dict(loss_dict):
|
||||
"""
|
||||
Reduce the loss dictionary from all processes so that process with rank
|
||||
0 has the averaged results. Returns a dict with the same fields as
|
||||
loss_dict, after reduction.
|
||||
"""
|
||||
world_size = dist_util.get_world_size()
|
||||
if world_size < 2:
|
||||
return loss_dict
|
||||
with torch.no_grad():
|
||||
loss_names = []
|
||||
all_losses = []
|
||||
for k in sorted(loss_dict.keys()):
|
||||
loss_names.append(k)
|
||||
all_losses.append(loss_dict[k])
|
||||
all_losses = torch.stack(all_losses, dim=0)
|
||||
dist.reduce(all_losses, dst=0)
|
||||
if dist.get_rank() == 0:
|
||||
# only main process gets accumulated, so only divide by
|
||||
# world_size in this case
|
||||
all_losses /= world_size
|
||||
reduced_losses = {k: v for k, v in zip(loss_names, all_losses)}
|
||||
return reduced_losses
|
||||
|
||||
|
||||
def do_train(cfg, model,
|
||||
data_loader,
|
||||
optimizer,
|
||||
scheduler,
|
||||
checkpointer,
|
||||
device,
|
||||
arguments,
|
||||
args):
|
||||
logger = logging.getLogger("SSD.trainer")
|
||||
logger.info("Start training ...")
|
||||
meters = MetricLogger()
|
||||
|
||||
model.train()
|
||||
save_to_disk = dist_util.get_rank() == 0
|
||||
if args.use_tensorboard and save_to_disk:
|
||||
import tensorboardX
|
||||
|
||||
summary_writer = tensorboardX.SummaryWriter(log_dir=os.path.join(cfg.OUTPUT_DIR, 'tf_logs'))
|
||||
else:
|
||||
summary_writer = None
|
||||
|
||||
max_iter = len(data_loader)
|
||||
start_iter = arguments["iteration"]
|
||||
start_training_time = time.time()
|
||||
end = time.time()
|
||||
for iteration, (images, targets, _) in enumerate(data_loader, start_iter):
|
||||
iteration = iteration + 1
|
||||
arguments["iteration"] = iteration
|
||||
scheduler.step()
|
||||
|
||||
images = images.to(device)
|
||||
targets = targets.to(device)
|
||||
loss_dict = model(images, targets=targets)
|
||||
loss = sum(loss for loss in loss_dict.values())
|
||||
|
||||
# reduce losses over all GPUs for logging purposes
|
||||
loss_dict_reduced = reduce_loss_dict(loss_dict)
|
||||
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
|
||||
meters.update(total_loss=losses_reduced, **loss_dict_reduced)
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
batch_time = time.time() - end
|
||||
end = time.time()
|
||||
meters.update(time=batch_time)
|
||||
if iteration % args.log_step == 0:
|
||||
eta_seconds = meters.time.global_avg * (max_iter - iteration)
|
||||
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
|
||||
logger.info(
|
||||
meters.delimiter.join([
|
||||
"iter: {iter:06d}",
|
||||
"lr: {lr:.5f}",
|
||||
'{meters}',
|
||||
"eta: {eta}",
|
||||
'mem: {mem}M',
|
||||
]).format(
|
||||
iter=iteration,
|
||||
lr=optimizer.param_groups[0]['lr'],
|
||||
meters=str(meters),
|
||||
eta=eta_string,
|
||||
mem=round(torch.cuda.max_memory_allocated() / 1024.0 / 1024.0),
|
||||
)
|
||||
)
|
||||
if summary_writer:
|
||||
global_step = iteration
|
||||
summary_writer.add_scalar('losses/total_loss', losses_reduced, global_step=global_step)
|
||||
for loss_name, loss_item in loss_dict_reduced.items():
|
||||
summary_writer.add_scalar('losses/{}'.format(loss_name), loss_item, global_step=global_step)
|
||||
summary_writer.add_scalar('lr', optimizer.param_groups[0]['lr'], global_step=global_step)
|
||||
|
||||
if iteration % args.save_step == 0:
|
||||
checkpointer.save("model_{:06d}".format(iteration), **arguments)
|
||||
|
||||
if args.eval_step > 0 and iteration % args.eval_step == 0 and not iteration == max_iter:
|
||||
eval_results = do_evaluation(cfg, model, distributed=args.distributed, iteration=iteration)
|
||||
if dist_util.get_rank() == 0 and summary_writer:
|
||||
for eval_result, dataset in zip(eval_results, cfg.DATASETS.TEST):
|
||||
write_metric(eval_result['metrics'], 'metrics/' + dataset, summary_writer, iteration)
|
||||
model.train() # *IMPORTANT*: change to train mode after eval.
|
||||
|
||||
checkpointer.save("model_final", **arguments)
|
||||
# compute training time
|
||||
total_training_time = int(time.time() - start_training_time)
|
||||
total_time_str = str(datetime.timedelta(seconds=total_training_time))
|
||||
logger.info("Total training time: {} ({:.4f} s / it)".format(total_time_str, total_training_time / max_iter))
|
||||
return model
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user