def train(**kwargs): """训练过程""" # 加载配置文件中的各种参数设置 OPT._parse(kwargs) # 数据集 dataset = Dataset(opt=OPT) print("加载数据集") dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=True, num_workers=OPT.num_workers) # 测试集 testset = TestDataset(opt=OPT) test_dataloader = DataLoader(dataset=testset, batch_size=1, shuffle=False, num_workers=OPT.num_workers, pin_memory=True) # 模型 faster_rcnn = FasterRCNNVGG16() print("模型加载完成") trainer = FasterRCNNTrainer(faster_rcnn).cuda() best_map = 0 # 最好的map lr_ = OPT.lr # 学习率 for epoch in range(OPT.epoch): print("Epoch: %s/%s" % (epoch, OPT.epoch - 1)) print("-" * 10) trainer.reset_meters() # 每次epoch的开始将损失函数清零 for ii, (img, bbox_, label_, scale) in pb.progressbar(enumerate(dataloader), max_value=len(dataloader)): scale = scalar(scale) # 原图和处理后的图片之间的一个缩放比例 img, bbox, label = img.cuda(), bbox_.cuda(), label_.cuda() trainer.train_step(imgs=img, bboxes=bbox, labels=label, scale=scale) print("train:", trainer.get_meter_data()) # if (ii + 1) % OPT.plot_every == 0: # print(trainer.get_meter_data()) trainer.eval() for jj, (img, size, _, bbox, label, _) in pb.progressbar(enumerate(test_dataloader), max_value=len(test_dataloader)): img, bbox, label = img.cuda(), bbox.cuda(), label.cuda() trainer.val_step(img, size, bbox, label) print("val:", trainer.get_meter_data()) eval_result = evaluate(dataloader=test_dataloader, faster_rcnn=faster_rcnn, test_num=OPT.test_num) print("mAP: %.4f" % eval_result["mAP"]) print() trainer.train()
def train(**kwargs): opt._parse(kwargs) image_folder_path = 'DataSets/images/' cvs_file_path = 'DataSets/labels.csv' dataset = DataSets(cvs_file_path, image_folder_path) data_size = len(dataset) indices = list(range(data_size)) split = int(np.floor(data_size * 0.2)) np.random.seed(42) np.random.shuffle(indices) train_indices, val_indices = indices[split:], indices[:split] train_sampler = torch.utils.data.SubsetRandomSampler(train_indices) valid_sampler = torch.utils.data.SubsetRandomSampler(val_indices) train_loader = torch.utils.data.DataLoader(dataset, batch_size=1, sampler=train_sampler) val_loader = torch.utils.data.DataLoader(dataset, batch_size=1, sampler=valid_sampler) print('load data') avg_loss = AverageValueMeter() ma20_loss = MovingAverageValueMeter(windowsize=20) faster_rcnn = FasterRCNNVGG16() print('model construct completed') start_epoch = 0 best_map = -100 trainer = FasterRCNNTrainer(faster_rcnn).cuda() optimizer = optim.SGD(trainer.faster_rcnn.parameters(), lr=opt.lr, momentum=0.9) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) if opt.load_path: print('load pretrained model from %s' % opt.load_path) checkpoint = torch.load(opt.load_path) start_epoch = checkpoint['epoch'] best_map = checkpoint['best_map'] trainer.faster_rcnn.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) print("> Loaded checkpoint '{}' (epoch {})".format( args.resume, start_epoch)) #trainer.vis.text(dataset.db.label_names, win='labels') # set tensor-board for visualization writer = SummaryWriter('runs/' + opt.log_root) for epoch in range(start_epoch, opt.epoch): trainer.train(mode=True) #must set as that in tranning for ii, (img, _, _, bbox_, label_, scale, _) in enumerate(train_loader): scale = at.scalar(scale) img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda() optimizer.zero_grad() loss = trainer.forward(img, bbox, label, scale) loss.total_loss.backward() optimizer.step() #print(loss) #print(loss.total_loss) loss_value = loss.total_loss.cpu().data.numpy() avg_loss.add(float(loss_value)) ma20_loss.add(float(loss_value)) print( '[epoch:{}/{}] [batch:{}/{}] [sample_loss:{:.4f}] [avg_loss:{:.4f}] [ma20_loss:{:.4f}]' .format(epoch, opt.epoch, ii + 1, len(train_loader), loss.total_loss.data, avg_loss.value()[0], ma20_loss.value()[0])) if (ii + 1) % opt.plot_every == 0: niter = epoch * len(train_loader) + ii writer.add_scalar('Train/Loss', ma20_loss.value()[0], niter) eval_result = eval(val_loader, faster_rcnn, test_num=opt.test_num) print(eval_result['map']) if eval_result['map'] > best_map: best_map = eval_result['map'] state = { "epoch": epoch + 1, "best_map": best_map, "model_state": trainer.faster_rcnn.state_dict(), "optimizer_state": optimizer.state_dict() } torch.save(state, opt.model_para) scheduler.step() state = { "epoch": epoch + 1, "best_map": best_map, "model_state": trainer.faster_rcnn.state_dict(), "optimizer_state": optimizer.state_dict() } torch.save(state, 'last_epoch.pkl') writer.close()