282 lines
11 KiB
Python
282 lines
11 KiB
Python
|
from __future__ import division
|
|||
|
|
|
|||
|
|
from models import *
|
|||
|
|
from utils.logger import *
|
|||
|
|
from utils.utils import *
|
|||
|
|
from utils.datasets import *
|
|||
|
|
from utils.parse_config import *
|
|||
|
|
from utils.prune_utils import *
|
|||
|
|
from test import evaluate
|
|||
|
|
|
|||
|
|
# 调试用的模块，reload用于代码热重载
|
|||
|
|
from importlib import reload
|
|||
|
|
import debug_utils
|
|||
|
|
|
|||
|
|
from terminaltables import AsciiTable
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import datetime
|
|||
|
|
import argparse
|
|||
|
|
|
|||
|
|
# import logging
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
from torch.utils.data import DataLoader
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--epochs", type=int, default=100, help="number of epochs")
|
|||
|
|
parser.add_argument("--batch_size", type=int, default=56, help="size of each image batch") # default=16
|
|||
|
|
parser.add_argument("--model_def", type=str, default="config/yolov3-person.cfg", help="path to model definition file") #default="config/yolov3-hand.cfg"
|
|||
|
|
parser.add_argument("--data_config", type=str, default="config/person.data", help="path to data config file") #default="config/oxfordhand.data"
|
|||
|
|
# parser.add_argument("--pretrained_weights", type=str, default="weights/darknet53.conv.74",
|
|||
|
|
parser.add_argument("--pretrained_weights", '-pre', type=str,
|
|||
|
|
default="weights/yolov3.weights", help="if specified starts from checkpoint model")
|
|||
|
|
parser.add_argument("--n_cpu", type=int, default=4, help="number of cpu threads to use during batch generation")
|
|||
|
|
parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
|
|||
|
|
parser.add_argument("--checkpoint_interval", type=int, default=20, help="interval between saving model weights") #default=5
|
|||
|
|
parser.add_argument("--evaluation_interval", type=int, default=1, help="interval evaluations on validation set")
|
|||
|
|
parser.add_argument("--multiscale_training", default=False, help="allow for multi-scale training")
|
|||
|
|
parser.add_argument("--debug_file", type=str, default="debug", help="enter ipdb if dir exists")
|
|||
|
|
|
|||
|
|
parser.add_argument('--learning_rate', '-lr', dest='lr', type=float, default=1e-3, help='initial learning rate')
|
|||
|
|
|
|||
|
|
parser.add_argument('--sparsity-regularization', '-sr', dest='sr', action='store_true',
|
|||
|
|
help='train with channel sparsity regularization')
|
|||
|
|
parser.add_argument('--s', type=float, default=0.01, help='scale sparse rate')
|
|||
|
|
|
|||
|
|
opt = parser.parse_args()
|
|||
|
|
print(opt)
|
|||
|
|
|
|||
|
|
logger = Logger("logs")
|
|||
|
|
|
|||
|
|
# 设置随机数种子
|
|||
|
|
init_seeds()
|
|||
|
|
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
|
|||
|
|
timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%m%d%H%M')
|
|||
|
|
|
|||
|
|
os.makedirs("output", exist_ok=True)
|
|||
|
|
os.makedirs("checkpoints", exist_ok=True)
|
|||
|
|
|
|||
|
|
# Get data configuration
|
|||
|
|
data_config = parse_data_config(opt.data_config)
|
|||
|
|
train_path = data_config["train"]
|
|||
|
|
valid_path = data_config["valid"]
|
|||
|
|
class_names = load_classes(data_config["names"])
|
|||
|
|
|
|||
|
|
# Initiate model
|
|||
|
|
model = Darknet(opt.model_def).to(device)
|
|||
|
|
model.apply(weights_init_normal)
|
|||
|
|
|
|||
|
|
#打印模型
|
|||
|
|
# print(model)
|
|||
|
|
|
|||
|
|
# If specified we start from checkpoint
|
|||
|
|
#是否加载预训练模型
|
|||
|
|
# if opt.pretrained_weights:
|
|||
|
|
# if opt.pretrained_weights.endswith(".pth"):
|
|||
|
|
# model.load_state_dict(torch.load(opt.pretrained_weights))
|
|||
|
|
# else:
|
|||
|
|
# model.load_darknet_weights(opt.pretrained_weights)
|
|||
|
|
|
|||
|
|
_, _, prune_idx= parse_module_defs(model.module_defs) # 剪枝需要的参数S
|
|||
|
|
# prune_idx = \
|
|||
|
|
# [0,
|
|||
|
|
# 2,
|
|||
|
|
# 6, 9,
|
|||
|
|
# 13, 16, 19, 22, 25, 28, 31, 34,
|
|||
|
|
# 38, 41, 44, 47, 50, 53, 56, 59,
|
|||
|
|
# 63, 66, 69, 72,
|
|||
|
|
# 75, 76, 77, 78, 79, 80,
|
|||
|
|
# #84,
|
|||
|
|
# 87, 88, 89, 90, 91, 92,
|
|||
|
|
# #96,
|
|||
|
|
# 99,100,101,102,103, 104]
|
|||
|
|
|
|||
|
|
# Get dataloader 导入数据
|
|||
|
|
dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training)
|
|||
|
|
dataloader = torch.utils.data.DataLoader(
|
|||
|
|
dataset,
|
|||
|
|
batch_size=opt.batch_size,
|
|||
|
|
shuffle=True,
|
|||
|
|
num_workers=opt.n_cpu,
|
|||
|
|
pin_memory=True,
|
|||
|
|
collate_fn=dataset.collate_fn
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9)
|
|||
|
|
|
|||
|
|
metrics = [
|
|||
|
|
"grid_size",
|
|||
|
|
"loss",
|
|||
|
|
"x", "y", "w", "h",
|
|||
|
|
"conf",
|
|||
|
|
"cls", "cls_acc",
|
|||
|
|
"recall50", "recall75",
|
|||
|
|
"precision",
|
|||
|
|
"conf_obj", "conf_noobj",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 配置日志记录器
|
|||
|
|
# logger = logging.getLogger('my_logger')
|
|||
|
|
# logger.setLevel(logging.INFO)
|
|||
|
|
# handler = logging.FileHandler('log.txt')
|
|||
|
|
# handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'))
|
|||
|
|
# logger.addHandler(handler)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# TFLOPS = 0
|
|||
|
|
tmp_loss = 0
|
|||
|
|
min_loss = 100
|
|||
|
|
best_mAP = 0
|
|||
|
|
|
|||
|
|
for epoch in range(opt.epochs):
|
|||
|
|
|
|||
|
|
# 进入调试模式
|
|||
|
|
if os.path.exists(opt.debug_file):
|
|||
|
|
import ipdb
|
|||
|
|
ipdb.set_trace()
|
|||
|
|
|
|||
|
|
model.train()
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
sr_flag = get_sr_flag(epoch, opt.sr)
|
|||
|
|
|
|||
|
|
obtain_num_parameters = lambda model:sum([param.nelement() for param in model.parameters()])
|
|||
|
|
origin_nparameters = obtain_num_parameters(model)
|
|||
|
|
print("Parameters : ", f"{origin_nparameters}")
|
|||
|
|
print("Model size:", sum([p.numel() for name, p in model.named_parameters()])) # 输出模型大小
|
|||
|
|
|
|||
|
|
for batch_i, (_, imgs, targets) in enumerate(dataloader):
|
|||
|
|
batches_done = len(dataloader) * epoch + batch_i
|
|||
|
|
|
|||
|
|
imgs = imgs.to(device)
|
|||
|
|
targets = targets.to(device)
|
|||
|
|
|
|||
|
|
loss, outputs = model(imgs, targets)
|
|||
|
|
|
|||
|
|
optimizer.zero_grad()
|
|||
|
|
loss.backward()
|
|||
|
|
|
|||
|
|
BNOptimizer.updateBN(sr_flag, model.module_list, opt.s, prune_idx)
|
|||
|
|
|
|||
|
|
optimizer.step()
|
|||
|
|
|
|||
|
|
# acc = (outputs > 0.5).float().mean()
|
|||
|
|
# logger.info('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Acc: {:.4f}'.format(epoch, epoch, i+1, len(dataloader), loss.item(), acc.item()))
|
|||
|
|
|
|||
|
|
# ----------------
|
|||
|
|
# Log progress
|
|||
|
|
# ----------------
|
|||
|
|
|
|||
|
|
log_str = "\n---- [Epoch %d/%d, Batch %d/%d] ----\n" % (epoch, opt.epochs, batch_i, len(dataloader))
|
|||
|
|
|
|||
|
|
metric_table = [["Metrics", *[f"YOLO Layer {i}" for i in range(len(model.yolo_layers))]]]
|
|||
|
|
|
|||
|
|
# Log metrics at each YOLO layer
|
|||
|
|
formats = {m: "%.6f" for m in metrics}
|
|||
|
|
formats["grid_size"] = "%2d"
|
|||
|
|
formats["cls_acc"] = "%.2f%%"
|
|||
|
|
for metric in metrics:
|
|||
|
|
row_metrics = [formats[metric] % yolo.metrics.get(metric, 0) for yolo in model.yolo_layers]
|
|||
|
|
metric_table += [[metric, *row_metrics]]
|
|||
|
|
|
|||
|
|
log_str += AsciiTable(metric_table).table
|
|||
|
|
log_str += f"\nTotal loss {loss.item()}"
|
|||
|
|
|
|||
|
|
#记录每个batch的最后一次loss
|
|||
|
|
tmp_loss = loss.item()
|
|||
|
|
|
|||
|
|
# Determine approximate time left for epoch
|
|||
|
|
epoch_batches_left = len(dataloader) - (batch_i + 1)
|
|||
|
|
time_left = datetime.timedelta(seconds=epoch_batches_left * (time.time() - start_time) / (batch_i + 1))
|
|||
|
|
log_str += f"\n---- ETA {time_left}"
|
|||
|
|
|
|||
|
|
print(log_str)
|
|||
|
|
|
|||
|
|
# Tensorboard logging
|
|||
|
|
tensorboard_log = []
|
|||
|
|
for i, yolo in enumerate(model.yolo_layers):
|
|||
|
|
for name, metric in yolo.metrics.items():
|
|||
|
|
# 选择部分指标写入tensorboard
|
|||
|
|
if name not in {"grid_size", "x", "y", "w", "h", "cls_acc"}:
|
|||
|
|
tensorboard_log += [(f"{name}_{i+1}", metric)]
|
|||
|
|
tensorboard_log += [("loss", loss.item())]
|
|||
|
|
tensorboard_log += [("lr", optimizer.param_groups[0]['lr'])]
|
|||
|
|
logger.list_of_scalars_summary('train', tensorboard_log, batches_done)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
if epoch % opt.evaluation_interval == 0:
|
|||
|
|
print("\n---- Evaluating Model ----")
|
|||
|
|
# Evaluate the model on the validation set
|
|||
|
|
precision, recall, AP, f1, ap_class = evaluate(
|
|||
|
|
model,
|
|||
|
|
path=valid_path,
|
|||
|
|
iou_thres=0.5,
|
|||
|
|
conf_thres=0.01,
|
|||
|
|
nms_thres=0.5,
|
|||
|
|
img_size=opt.img_size,
|
|||
|
|
batch_size=8,
|
|||
|
|
)
|
|||
|
|
evaluation_metrics = [
|
|||
|
|
("val_precision", precision.mean()),
|
|||
|
|
("val_recall", recall.mean()),
|
|||
|
|
("val_mAP", AP.mean()),
|
|||
|
|
("val_f1", f1.mean()),
|
|||
|
|
]
|
|||
|
|
logger.list_of_scalars_summary('valid', evaluation_metrics, epoch)
|
|||
|
|
|
|||
|
|
# Print class APs and mAP
|
|||
|
|
ap_table = [["Index", "Class name", "AP"]]
|
|||
|
|
for i, c in enumerate(ap_class):
|
|||
|
|
ap_table += [[c, class_names[c], "%.5f" % AP[i]]]
|
|||
|
|
print(AsciiTable(ap_table).table)
|
|||
|
|
print(f"---- mAP {AP.mean()}")
|
|||
|
|
|
|||
|
|
# 往tensorboard中记录bn权重分布
|
|||
|
|
bn_weights = gather_bn_weights(model.module_list, prune_idx)
|
|||
|
|
logger.writer.add_histogram('bn_weights/hist', bn_weights.numpy(), epoch, bins='doane')
|
|||
|
|
|
|||
|
|
if(best_mAP < AP.mean()):
|
|||
|
|
best_mAP = AP.mean()
|
|||
|
|
torch.save(model.state_dict(), f"checkpoints/xs_nopretrain_pt_best_mAP_{epoch}_yolov3_ckpt.pth") # 保存mAP最高的模型
|
|||
|
|
if(tmp_loss < min_loss):
|
|||
|
|
min_loss = tmp_loss
|
|||
|
|
torch.save(model.state_dict(), f"checkpoints/xs_nopretrain_pt_min_loss_{epoch}_yolov3_ckpt.pth") # 保存loss最低的模型
|
|||
|
|
|
|||
|
|
if epoch % opt.checkpoint_interval == 0 or epoch == opt.epochs - 1:
|
|||
|
|
|
|||
|
|
torch.save(model.state_dict(), f"checkpoints/xs_nopretrain_yolov3_ckpt_{epoch}_{timestamp}.pth")
|
|||
|
|
print(f"model has been saved:checkpoints/xs_nopretrain_yolov3_ckpt_{epoch}_{timestamp}.pth")
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# obtain_num_parameters = lambda model:sum([param.nelement() for param in model.parameters()])
|
|||
|
|
# origin_nparameters = obtain_num_parameters(model)
|
|||
|
|
# print("Parameters : ", f"{origin_nparameters}")
|
|||
|
|
|
|||
|
|
# GPUs = torch.hub.list_hub_servers()
|
|||
|
|
# flops = []
|
|||
|
|
# parameters = []
|
|||
|
|
|
|||
|
|
# for server in GPUs:
|
|||
|
|
# # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
# flops.append(torch.Tensor(np.sum(np.square(model.parameters()[:3])).item()))
|
|||
|
|
# parameters.append(torch.Tensor(model.parameters()))
|
|||
|
|
# model.to(device)
|
|||
|
|
# flops = np.array(flops)
|
|||
|
|
# flops = flops.reshape(-1, flops.shape[-1])
|
|||
|
|
# TFLOPS += flops.item()
|
|||
|
|
# parameters = np.array(parameters)
|
|||
|
|
# parameters = parameters.reshape(-1, parameters.shape[-1])
|
|||
|
|
# print('GPUs:', GPUs)
|
|||
|
|
# print('FLOPS:', flops.item())
|
|||
|
|
# print('Params:', parameters.item())
|
|||
|
|
|
|||
|
|
# print('TFLOPS:', TFLOPS)
|