def main(args):
init_distributed_mode(args)
print(args)
device = torch.device(args.device)
# 用来保存coco_info的文件
results_file = "results{}.txt".format(
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
# Data loading code
print("Loading data")
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
COCO_root = args.data_path
# load train data set
# coco2017 -> annotations -> instances_train2017.json
train_data_set = CocoDetection(COCO_root, "train", data_transform["train"])
# load validation data set
# coco2017 -> annotations -> instances_val2017.json
val_data_set = CocoDetection(COCO_root, "val", data_transform["val"])
print("Creating data loaders")
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_data_set)
test_sampler = torch.utils.data.distributed.DistributedSampler(
val_data_set)
else:
train_sampler = torch.utils.data.RandomSampler(train_data_set)
test_sampler = torch.utils.data.SequentialSampler(val_data_set)
if args.aspect_ratio_group_factor >= 0:
# 统计所有图像比例在bins区间中的位置索引
group_ids = create_aspect_ratio_groups(
train_data_set, k=args.aspect_ratio_group_factor)
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids,
args.batch_size)
else:
train_batch_sampler = torch.utils.data.BatchSampler(train_sampler,
args.batch_size,
drop_last=True)
data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_sampler=train_batch_sampler,
num_workers=args.workers,
collate_fn=train_data_set.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
val_data_set,
batch_size=1,
sampler=test_sampler,
num_workers=args.workers,
collate_fn=train_data_set.collate_fn)
print("Creating model")
# create model num_classes equal background + 80 classes
model = create_model(num_classes=args.num_classes + 1, device=device)
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
# 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练
if args.resume:
# If map_location is missing, torch.load will first load the module to CPU
# and then copy each parameter to where it was saved,
# which would result in all processes on the same machine using the same set of devices.
checkpoint = torch.load(
args.resume, map_location='cpu') # 读取之前保存的权重文件(包括优化器以及学习率策略)
model_without_ddp.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
train_loss = []
learning_rate = []
val_map = []
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
mean_loss, lr = utils.train_one_epoch(model,
optimizer,
data_loader,
device,
epoch,
args.print_freq,
warmup=True)
# update learning rate
lr_scheduler.step()
# evaluate after every epoch
coco_info = utils.evaluate(model, data_loader_test, device=device)
# 只在主进程上进行写操作
if args.rank in [-1, 0]:
train_loss.append(mean_loss.item())
learning_rate.append(lr)
val_map.append(coco_info[1]) # pascal mAP
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [
str(round(i, 4))
for i in coco_info + [mean_loss.item(), lr]
]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
if args.output_dir:
# 只在主节点上执行保存权重操作
save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'args': args,
'epoch': epoch
}, os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
if args.rank in [-1, 0]:
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
def main(args):
device = torch.device(args.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 用来保存coco_info的文件
results_file = "results{}.txt".format(
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
COCO_root = args.data_path
# load train data set
# coco2017 -> annotations -> instances_train2017.json
train_data_set = CocoDetection(COCO_root, "train", data_transform["train"])
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = args.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# load validation data set
# coco2017 -> annotations -> instances_val2017.json
val_data_set = CocoDetection(COCO_root, "val", data_transform["val"])
val_data_set_loader = torch.utils.data.DataLoader(
val_data_set,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# create model num_classes equal background + 80 classes
model = create_model(num_classes=args.num_classes + 1)
# print(model)
model.to(device)
train_loss = []
learning_rate = []
val_map = []
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
# 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练
if args.resume:
# If map_location is missing, torch.load will first load the module to CPU
# and then copy each parameter to where it was saved,
# which would result in all processes on the same machine using the same set of devices.
checkpoint = torch.load(
args.resume, map_location=device) # 读取之前保存的权重文件(包括优化器以及学习率策略)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch, printing every 50 iterations
mean_loss, lr = utils.train_one_epoch(model,
optimizer,
train_data_loader,
device,
epoch,
print_freq=50,
warmup=True)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [
str(round(i, 4)) for i in coco_info + [mean_loss.item(), lr]
]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
# save weights
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch
}
torch.save(save_files,
"./save_weights/mobile-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 用来保存coco_info的文件
results_file = "results{}.txt".format(
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists("save_weights"):
os.makedirs("save_weights")
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
COCO_root = "/data/coco2017"
# load train data set
# coco2017 -> annotations -> instances_train2017.json
train_data_set = CocoDetection(COCO_root, "train", data_transform["train"])
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = 8
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# load validation data set
# coco2017 -> annotations -> instances_val2017.json
val_data_set = CocoDetection(COCO_root, "val", data_transform["val"])
val_data_set_loader = torch.utils.data.DataLoader(
val_data_set,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# create model num_classes equal background + 80 classes
model = create_model(num_classes=81)
# print(model)
model.to(device)
train_loss = []
learning_rate = []
val_map = []
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.001,
momentum=0.9,
weight_decay=1e-4)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[16, 22],
gamma=0.1)
num_epochs = 26
for epoch in range(num_epochs):
# train for one epoch, printing every 50 iterations
mean_loss, lr = utils.train_one_epoch(model,
optimizer,
train_data_loader,
device,
epoch,
print_freq=50,
warmup=True)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [
str(round(i, 4)) for i in coco_info + [mean_loss.item(), lr]
]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
# save weights
# 仅保存最后5个epoch的权重
if epoch in range(num_epochs)[-5:]:
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch
}
torch.save(save_files,
"./save_weights/mobile-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
def main(parser_data):
device = torch.device(
parser_data.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
data_transform = {"val": transforms.Compose([transforms.ToTensor()])}
# read class_indict
label_json_path = './coco80_indices.json'
assert os.path.exists(
label_json_path), "json file {} dose not exist.".format(
label_json_path)
json_file = open(label_json_path, 'r')
category_index = json.load(json_file)
coco_root = parser_data.data_path
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = parser_data.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
# load validation data set
val_dataset = CocoDetection(coco_root, "val", data_transform["val"])
val_dataset_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=val_dataset.collate_fn)
# create model
vgg_feature = vgg(model_name="vgg16",
weights_path="./backbone/vgg16.pth").features
backbone = torch.nn.Sequential(*list(
vgg_feature._modules.values())[:-1]) # 删除feature中最后的maxpool层
backbone.out_channels = 512
anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'], # 在哪些特征层上进行roi pooling
output_size=[7, 7], # roi_pooling输出特征矩阵尺寸
sampling_ratio=2) # 采样率
# num_classes equal 80 + background classes
model = FasterRCNN(backbone=backbone,
num_classes=parser_data.num_classes + 1,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# 载入你自己训练好的模型权重
weights_path = parser_data.weights
assert os.path.exists(weights_path), "not found {} file.".format(
weights_path)
weights_dict = torch.load(weights_path, map_location=device)
model.load_state_dict(weights_dict['model'])
# print(model)
model.to(device)
# evaluate on the val dataset
cpu_device = torch.device("cpu")
coco91to80 = val_dataset.coco91to80
coco80to91 = dict([(str(v), k) for k, v in coco91to80.items()])
results = []
model.eval()
with torch.no_grad():
for image, targets in tqdm(val_dataset_loader, desc="validation..."):
# 将图片传入指定设备device
image = list(img.to(device) for img in image)
# inference
outputs = model(image)
outputs = [{k: v.to(cpu_device)
for k, v in t.items()} for t in outputs]
# 遍历每张图像的预测结果
for target, output in zip(targets, outputs):
if len(output) == 0:
continue
img_id = int(target["image_id"])
per_image_boxes = output["boxes"]
# 对于coco_eval, 需要的每个box的数据格式为[x_min, y_min, w, h]
# 而我们预测的box格式是[x_min, y_min, x_max, y_max]，所以需要转下格式
per_image_boxes[:, 2:] -= per_image_boxes[:, :2]
per_image_classes = output["labels"]
per_image_scores = output["scores"]
# 遍历每个目标的信息
for object_score, object_class, object_box in zip(
per_image_scores, per_image_classes, per_image_boxes):
object_score = float(object_score)
# 要将类别信息还原回coco91中
coco80_class = int(object_class)
coco91_class = int(coco80to91[str(coco80_class)])
# We recommend rounding coordinates to the nearest tenth of a pixel
# to reduce resulting JSON file size.
object_box = [round(b, 2) for b in object_box.tolist()]
res = {
"image_id": img_id,
"category_id": coco91_class,
"bbox": object_box,
"score": round(object_score, 3)
}
results.append(res)
# accumulate predictions from all images
# write predict results into json file
json_str = json.dumps(results, indent=4)
with open('predict_tmp.json', 'w') as json_file:
json_file.write(json_str)
# accumulate predictions from all images
coco_true = val_dataset.coco
coco_pre = coco_true.loadRes('predict_tmp.json')
coco_evaluator = COCOeval(cocoGt=coco_true,
cocoDt=coco_pre,
iouType="bbox")
coco_evaluator.evaluate()
coco_evaluator.accumulate()
coco_evaluator.summarize()
# calculate COCO info for all classes
coco_stats, print_coco = summarize(coco_evaluator)
# calculate voc info for every classes(IoU=0.5)
voc_map_info_list = []
for i in range(len(category_index)):
stats, _ = summarize(coco_evaluator, catId=i)
voc_map_info_list.append(" {:15}: {}".format(
category_index[str(i + 1)], stats[1]))
print_voc = "\n".join(voc_map_info_list)
print(print_voc)
# 将验证结果保存至txt文件中
with open("record_mAP.txt", "w") as f:
record_lines = [
"COCO results:", print_coco, "", "mAP(IoU=0.5) for each category:",
print_voc
]
f.write("\n".join(record_lines))