def coco_evaluate(model, data_loader, device):
n_threads = torch.get_num_threads()
# FIXME remove this and make paste_masks_in_image run on the GPU
torch.set_num_threads(1)
cpu_device = torch.device("cpu")
model.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
header = 'Test:'
coco = get_coco_api_from_dataset(data_loader.dataset)
iou_types = _get_iou_types(model)
coco_evaluator = CocoEvaluator(coco, iou_types)
for image, targets in metric_logger.log_every(data_loader, 100, header):
image = list(img.to(device) for img in image)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
torch.cuda.synchronize()
model_time = time.time()
outputs = model(image)
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
model_time = time.time() - model_time
res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_time = time.time() - evaluator_time
metric_logger.update(model_time=model_time, evaluator_time=evaluator_time)
# gather the stats from all processes
metric_logger.synchronize_between_processes()
print("Averaged stats:", metric_logger)
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
torch.set_num_threads(n_threads)
return coco_evaluator
def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq):
model.train()
metric_logger = utils.MetricLogger(delimiter=" ")
metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
header = 'Epoch: [{}]'.format(epoch)
lr_scheduler = None
if epoch == 0:
warmup_factor = 1. / 1000
warmup_iters = min(1000, len(data_loader) - 1)
lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)
for images, targets in metric_logger.log_every(data_loader, print_freq, header):
images = list(image.to(device) for image in images)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
# reduce losses over all GPUs for logging purposes
loss_dict_reduced = utils.reduce_dict(loss_dict)
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
loss_value = losses_reduced.item()
if not math.isfinite(loss_value):
print("Loss is {}, stopping training".format(loss_value))
print(loss_dict_reduced)
sys.exit(1)
optimizer.zero_grad()
losses.backward()
optimizer.step()
if lr_scheduler is not None:
lr_scheduler.step()
metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
metric_logger.update(lr=optimizer.param_groups[0]["lr"])
def voc_evaluate(model, data_loader, device):
n_threads = torch.get_num_threads()
torch.set_num_threads(1)
cpu_device = torch.device("cpu")
model.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
header = 'Test:'
all_boxes = [[] for i in range(21)]
image_index = []
for image, targets in metric_logger.log_every(data_loader, 100, header):
image = list(img.to(device) for img in image)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
torch.cuda.synchronize()
model_time = time.time()
outputs = model(image)
name = ''.join([chr(i) for i in targets[0]['name'].tolist()])
image_index.append(name)
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
image_boxes = [[] for i in range(21)]
for o in outputs:
for i in range(o['boxes'].shape[0]):
image_boxes[o['labels'][i]].extend([
torch.cat([o['boxes'][i], o['scores'][i].unsqueeze(0)], dim=0)
])
# makes sure that the all_boxes is filled with empty array when
# there are no boxes in image_boxes
for i in range(21):
if image_boxes[i] != []:
all_boxes[i].append([torch.stack(image_boxes[i])])
else:
all_boxes[i].append([])
model_time = time.time() - model_time
metric_logger.synchronize_between_processes()
all_boxes_gathered = utils.all_gather(all_boxes)
image_index_gathered = utils.all_gather(image_index)
# results from all processes are gathered here
if utils.is_main_process():
all_boxes = [[] for i in range(21)]
for abgs in all_boxes_gathered:
for ab, abg in zip(all_boxes, abgs):
ab += abg
image_index = []
for iig in image_index_gathered:
image_index += iig
_write_voc_results_file(all_boxes, image_index, data_loader.dataset.root,
data_loader.dataset._transforms.transforms[0].CLASSES)
_do_python_eval(data_loader)
torch.set_num_threads(n_threads)
def evaluate(model, data_loader, coco=None, device=None):
n_threads = torch.get_num_threads()
# FIXME remove this and make paste_masks_in_image run on the GPU
torch.set_num_threads(1)
cpu_device = torch.device("cpu")
model.eval()
metric_logger = utils.MetricLogger(delimiter=" ")
header = "Test: "
if coco is None:
coco = get_coco_api_from_dataset(data_loader.dataset)
iou_types = _get_iou_types(model)
coco_evaluator = CocoEvaluator(coco, iou_types)
for imgs, targets, paths, _, img_index in metric_logger.log_every(data_loader, 100, header):
imgs = imgs.to(device).float() / 255.0 # uint8 to float32, 0 - 255 to 0.0 - 1.0
# targets = targets.to(device)
# 当使用CPU时，跳过GPU相关指令
if device != torch.device("cpu"):
torch.cuda.synchronize(device)
model_time = time.time()
pred = model(imgs)[0] # only get inference result
pred = non_max_suppression(pred, conf_thres=0.001, iou_thres=0.6, multi_label=False)
outputs = []
for index, p in enumerate(pred):
if p is None:
p = torch.empty((0, 6), device=cpu_device)
boxes = torch.empty((0, 4), device=cpu_device)
else:
# xmin, ymin, xmax, ymax
boxes = p[:, :4]
# 注意这里传入的boxes格式必须是xmin, ymin, xmax, ymax，且为绝对坐标
info = {"boxes": boxes.to(cpu_device),
"labels": p[:, 5].to(device=cpu_device, dtype=torch.int64),
"scores": p[:, 4].to(cpu_device)}
outputs.append(info)
model_time = time.time() - model_time
res = {img_id: output for img_id, output in zip(img_index, outputs)}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_time = time.time() - evaluator_time
metric_logger.update(model_time=model_time, evaluator_time=evaluator_time)
# gather the stats from all processes
metric_logger.synchronize_between_processes()
print("Averaged stats:", metric_logger)
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
torch.set_num_threads(n_threads)
result_info = coco_evaluator.coco_eval[iou_types[0]].stats
return result_info
def train_one_epoch(model, optimizer, data_loader, device, epoch,
print_freq, accumulate, img_size, batch_size,
grid_min, grid_max, gs,
multi_scale=False, warmup=False):
model.train()
metric_logger = utils.MetricLogger(delimiter=" ")
metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
header = 'Epoch: [{}]'.format(epoch)
lr_scheduler = None
if epoch == 0 and warmup is True: # 当训练第一轮（epoch=0）时，启用warmup训练方式，可理解为热身训练
warmup_factor = 5.0 / 1000
warmup_iters = min(1000, len(data_loader) - 1)
lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)
enable_amp = True if "cuda" in device.type else False
mloss = torch.zeros(4).to(device) # mean losses
now_lr = 0.
nb = len(data_loader) # number of batches
# imgs: [batch_size, 3, img_size, img_size]
# targets: [num_obj, 6] , that number 6 means -> (img_index, obj_index, x, y, w, h)
# paths: list of img path
for i, (imgs, targets, paths, _, _) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
# ni 统计从epoch0开始的所有batch数
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device).float() / 255.0 # uint8 to float32, 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
# Multi-Scale
if multi_scale:
# 每训练64张图片，就随机修改一次输入图片大小，
# 由于label已转为相对坐标，故缩放图片不影响label的值
if ni % accumulate == 0: # adjust img_size (67% - 150%) every 1 batch
# 在给定最大最小输入尺寸范围内随机选取一个size(size为32的整数倍)
img_size = random.randrange(grid_min, grid_max + 1) * gs
sf = img_size / max(imgs.shape[2:]) # scale factor
# 如果图片最大边长不等于img_size, 则缩放图片，并将长和宽调整到32的整数倍
if sf != 1:
# gs: (pixels) grid size
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to 32-multiple)
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# 混合精度训练上下文管理器，如果在CPU环境中不起任何作用
with torch.cuda.amp.autocast(enabled=enable_amp):
pred = model(imgs)
# loss
loss_dict = compute_loss(pred, targets, model)
losses = sum(loss for loss in loss_dict.values())
# reduce losses over all GPUs for logging purpose
loss_dict_reduced = utils.reduce_dict(loss_dict)
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
loss_items = torch.cat((loss_dict_reduced["box_loss"],
loss_dict_reduced["obj_loss"],
loss_dict_reduced["class_loss"],
losses_reduced)).detach()
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
loss_value = losses_reduced.item()
if not torch.isfinite(losses_reduced):
print('WARNING: non-finite loss, ending training ', loss_value)
print("training image path: {}".format(",".join(paths)))
sys.exit(1)
# 每训练64张图片更新一次权重
# backward
losses *= batch_size / 64 # scale loss
losses.backward()
# optimize
if ni % accumulate == 0:
optimizer.step()
optimizer.zero_grad()
metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
now_lr = optimizer.param_groups[0]["lr"]
metric_logger.update(lr=now_lr)
if lr_scheduler is not None: # 第一轮使用warmup训练方式
lr_scheduler.step()
return mloss, now_lr
l1_reg='bic')
labels = np.argmax(train_label, axis=1)
if args.weight == "instance_level":
contributions_shap = compare_shap_and_KG(shap_values_train,
labels,
dataset=data)
shap_coeff = reduce_shap(contributions_shap, is_exponential)
elif args.weight == "bbox_level":
shap_weights = get_bbox_weight(shap_values_train,
is_exponential,
dataset=data)
print("Shap computed")
print('Test loss: ', loss, '\tTest accuracy: ', accuracy)
if j < num_epochs_detection:
#Train detection
metric_logger = uti.MetricLogger(delimiter=" ")
metric_logger.add_meter(
'lr', uti.SmoothedValue(window_size=1, fmt='{value:.6f}'))
header = 'Epoch: [{}]'.format(j + 1)
index = 0
detector.train()
for images, targets in metric_logger.log_every(train_loader, 250,
header):
images = list(image.to('cuda') for image in images)
targets = [{k: v.to('cuda')
for k, v in t.items()} for t in targets]
#SHAP if necessary
#-------
if (j > 0 or args.resume) and args.weight == "bbox_level":
loss_dict = detector(
