def test():
val_dataset = build_dataset(data_dir=DATA_DIR,
anno_file=f'{DATA_DIR}/annotations_aug.json',
classnames=CLASSNAMES,
batch_size=1,
shuffle=False,
num_workers=8,
is_train=False)
faster_rcnn = FasterRCNN(n_class=len(CLASSNAMES)+1)
files = sorted(list(glob.glob(f'{save_checkpoint_path}/checkpoint*.pkl')))
f = files[-1]
faster_rcnn.load(f)
mAP = eval(val_dataset,faster_rcnn,0,is_display=True,is_save_result=False,score_thresh=0.5)
print(mAP)
def test():
img_arr = read_image('demo.jpg')
img = t.from_numpy(img_arr)[None]
faster_rcnn = FasterRCNN()
trainer = FasterRCNNTrainer(faster_rcnn).cuda()
trainer.load('weights/chainer_best_model_converted_to_pytorch_0.7053.pth')
opt.caffe_pretrain = True
_bboxes, _labels, _scores = trainer.faster_rcnn.predict(img, visualize=True)
vis_bbox(at.tonumpy(img[0]),
at.tonumpy(_bboxes[0]),
at.tonumpy(_labels[0]).reshape(-1),
at.tonumpy(_scores[0]).reshape(-1))
def train():
train_dataset = build_dataset(data_dir=DATA_DIR,
anno_file=f'{DATA_DIR}/annotations_aug.json',
classnames=CLASSNAMES,
batch_size=1,
shuffle=True,
num_workers=4,
is_train=True,
use_all=True)
val_dataset = build_dataset(data_dir=DATA_DIR,
anno_file=f'{DATA_DIR}/annotations_aug.json',
classnames=CLASSNAMES,
batch_size=1,
shuffle=False,
num_workers=4,
is_train=False,
use_all=True)
faster_rcnn = FasterRCNN(n_class = len(CLASSNAMES)+1)
optimizer = optim.SGD(faster_rcnn.parameters(),momentum=0.9,lr=0.001)
writer = SummaryWriter()
for epoch in range(EPOCHS):
faster_rcnn.train()
dataset_len = len(train_dataset)
for batch_idx,(images,boxes,labels,image_sizes,img_ids) in tqdm(enumerate(train_dataset)):
rpn_loc_loss,rpn_cls_loss,roi_loc_loss,roi_cls_loss,total_loss = faster_rcnn(images,boxes,labels)
optimizer.step(total_loss)
writer.add_scalar('rpn_cls_loss', rpn_cls_loss.item(), global_step=dataset_len*epoch+batch_idx)
writer.add_scalar('rpn_loc_loss', rpn_loc_loss.item(), global_step=dataset_len*epoch+batch_idx)
writer.add_scalar('roi_loc_loss', roi_loc_loss.item(), global_step=dataset_len*epoch+batch_idx)
writer.add_scalar('roi_cls_loss', roi_cls_loss.item(), global_step=dataset_len*epoch+batch_idx)
writer.add_scalar('total_loss', total_loss.item(), global_step=dataset_len*epoch+batch_idx)
if batch_idx % 10 == 0:
loss_str = '\nrpn_loc_loss: %s \nrpn_cls_loss: %s \nroi_loc_loss: %s \nroi_cls_loss: %s \ntotoal_loss: %s \n'
print(loss_str % (rpn_loc_loss.item(),rpn_cls_loss.item(),roi_loc_loss.item(),roi_cls_loss.item(),total_loss.item()))
mAP = eval(val_dataset,faster_rcnn,epoch)
writer.add_scalar('map', mAP, global_step=epoch)
os.makedirs(save_checkpoint_path,exist_ok=True)
faster_rcnn.save(f"{save_checkpoint_path}/checkpoint_{epoch}.pkl")
def run_images(img_dir,classnames,checkpoint_path,save_path):
test_dataset = build_testdataset(img_dir=img_dir)
faster_rcnn = FasterRCNN(n_class=len(classnames)+1)
faster_rcnn.load(checkpoint_path)
faster_rcnn.eval()
results = {}
for batch_idx,(images,image_sizes,img_ids) in tqdm(enumerate(test_dataset)):
result = faster_rcnn.predict(images,score_thresh=0.01)
for img_id,(pred_boxes,pred_scores,pred_labels) in zip(img_ids,result):
objects = []
for box,label,score in zip(pred_boxes.numpy(),pred_labels.numpy(),pred_scores.numpy()):
bbox = {"xmin":float(box[0]),"ymin":float(box[1]),"xmax":float(box[2]),"ymax":float(box[3])}
category = classnames[label-1]
score = float(score)
objects.append({"bbox":bbox,"category":category,"score":score})
results[img_id+".jpg"]=objects
os.makedirs(save_path,exist_ok=True)
json.dump(results,open(os.path.join(save_path,"test_result.json"),"w"))
dummy_search_box = np.array([0, 0, bound[0], bound[1]]).reshape(1, -1)
anchors = G.gen_region_anchors(raw_anchors,
dummy_search_box,
bound,
K=K,
size=out_size)[0]
print(anchors.shape)
img_files = [
'img00337.jpg', 'img00832.jpg', 'img00995.jpg', 'img01879.jpg',
'road.jpg'
]
# img_files=['road.jpg']
model_path = './ckpt/model_660000.pkl'
model = FasterRCNN(im_width, im_height, pretrained=False)
model.load_weights(model_path=model_path)
model.cuda()
for f in img_files:
img_name = f[:f.rfind('.')]
img_path = './images/{}'.format(f)
image = cv2.imread(img_path)
canvas_det = image.copy()
canvas_prpsl = image.copy()
cls_bboxes, xyscale, proposals = im_detect(image, anchors, model)
draw_boxes(canvas_det, cls_bboxes, xyscale)
draw_proposals(canvas_prpsl, proposals, xyscale)
cv2.imwrite('{}_det.jpg'.format(img_name), canvas_det)
def main():
writer = SummaryWriter(log_dir=checkpoint_dir)
# wrap model init and training in a main(), otherwise it will mess up with python multithreading when num_workers > 0
model = FasterRCNN().to(device)
# check grad flag
# for name, param in model.named_parameters():
# if param.requires_grad:
# print(name)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
#optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', patience=3,
verbose=True) # dynamic LR scheduler
# Resume traning
if args.resume:
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
start_epoch = checkpoint['epoch']
print("Resume training from epoch {}".format(start_epoch + 1))
else:
start_epoch = 0
# Training
for epoch in range(start_epoch, num_epochs):
print("Epoch {}/{}".format(epoch + 1, num_epochs))
# Train
model.train()
loss_1, loss_2, loss_3, loss_4, loss_5 = 0, 0, 0, 0, 0 # running sum of loss
for batch_idx, (images, gt_boxes,
gt_classes) in enumerate(trainloader):
images, gt_boxes, gt_classes = images.to(device), gt_boxes.to(
device), gt_classes.to(device)
optimizer.zero_grad() # reset gradient
rois, pred_rois_scores, pred_rois_coeffs, rpn_class_loss, rpn_bbox_loss, rcnn_class_loss, rcnn_bbox_loss, rcnn_loss = model(
images, gt_boxes, gt_classes) # forward step
# rois: N x R x 4
# pred_rois_scores: N x R x 21 (only used during eval/inference)
# pred_rois_coeffs: N x R x 21*4 (only used during eval/inference)
# *loss: scalar (only used during training)
# during training, only rcnn_loss is used for backward
# during evaluation or inference, more post-processing should be done including:
# 1. Apply 'pred_rois_coeffs' to 'rois' to get the final RoIs, e.g. rois = rpn.utils.bbox_transform(rois, pred_rois_coeffs)
# 2. Get RoI class and prediction confidence by argmax 'pred_rois_classes', e.g. pred_rois_classes = torch.argmax(pred_rois_scores, dim=2) # N x R
# 3. For evaluation, compute accuracy; for inference, plot
rcnn_loss.backward() # calculate gradients
# plot_grad_flow(model.named_parameters())
# print loss every X steps (i.e. average loss of X * cfg.BATCH_SIZE samples)
# tensorboard log
step_no = epoch * len(trainloader) + batch_idx
writer.add_scalars(
'all_loss', {
'rpn_class_loss': rpn_class_loss.item(),
'rpn_bbox_loss': rpn_bbox_loss.item(),
'rcnn_class_loss': rcnn_class_loss.item(),
'rcnn_bbox_loss': rcnn_bbox_loss.item(),
'rcnn_loss': rcnn_loss.item()
}, step_no)
# print
loss_1 += rpn_class_loss.item()
loss_2 += rpn_bbox_loss.item()
loss_3 += rcnn_class_loss.item()
loss_4 += rcnn_bbox_loss.item()
loss_5 += rcnn_loss.item()
if (batch_idx + 1) % loss_print_step == 1:
print("> Step {}/{}".format(batch_idx, len(trainloader)),
end=', ',
flush=True)
print(
"rpn_class_loss:{:3f}, rpn_bbox_loss:{:3f}, rcnn_class_loss:{:3f}, rcnn_bbox_loss:{:3f}, rcnn_loss:{:.3f}"
.format(loss_1 / (batch_idx + 1), loss_2 / (batch_idx + 1),
loss_3 / (batch_idx + 1), loss_4 / (batch_idx + 1),
loss_5 / (batch_idx + 1)))
optimizer.step() # update parameters
scheduler.step(loss_5 /
len(trainloader)) # update LR based on loss trend
# print averaged loss per entire epoch
print(
"> Entire epoch loss: rpn_class_loss:{:3f}, rpn_bbox_loss:{:3f}, rcnn_class_loss:{:3f}, rcnn_bbox_loss:{:3f}, rcnn_loss:{:.3f}"
.format(loss_1 / len(trainloader), loss_2 / len(trainloader),
loss_3 / len(trainloader), loss_4 / len(trainloader),
loss_5 / len(trainloader)))
# save model checkpoints
if (epoch + 1) % model_save_step == 1:
print("Saving model checkpoint at epoch {}/{}".format(
epoch + 1, num_epochs))
checkpoint = {
'epoch': epoch +
1, # epoch is completed, should start from epoch+1 later on
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict()
}
torch.save(checkpoint, checkpoint_path)
return keeps
# Define the recall and total dictionary for calculating mAP
#recall = {}
#total = {}
#for i in range(21):
# recall[i] = 0
# total[i] = 0
# Define the test loader: Only use toTensor to keep the original image
testset = VOCDetection(root='./datasets', year='2007', image_set='train', download=False, transform=transforms.ToTensor())
testloader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=1)
# Load the model
checkpoint = torch.load('rcnn.pth', map_location=torch.device('cpu'))
model = FasterRCNN()
model.eval() # Turn to eval mode
model.load_state_dict(checkpoint['state_dict'])
print("Start eval...")
for batch_idx, (images, gt_boxes, gt_classes, gt_boxes_, gt_classes_) in enumerate(testloader):
# Use for partial training
if batch_idx <= 147:
continue
# Add the normalization step
img = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])(images[0])
out = model(img.unsqueeze(0),None,None)
# Get the keep index for all the predicted bounding boxes
class TrainEngine(object):
def __init__(self):
self.batch_size = 4
cfg.PHASE = 'TRAIN'
cfg[cfg.PHASE].IMS_PER_BATCH = self.batch_size
self.stride = STRIDE[net_type]
cfg.STRIDE = self.stride
cfg.GAIN = 0.02
self.basic_size = cfg.BASIC_SIZE
self.ratios = cfg.RATIOS
self.scales = cfg.SCALES
self.backbone_pretrained = True
self.lr_mult = 0.5
self.decay_ratio = 0.1
self.K = len(self.ratios) * len(self.scales)
self.im_w, self.im_h = RESOLUTION[net_type]
self.model = None
def xy2wh(self, boxes):
x = boxes[:, 0]
y = boxes[:, 1]
w = boxes[:, 2] - x + 1
h = boxes[:, 3] - y + 1
return x, y, w, h
def get_param_groups(self, base_lr):
backbone_lr = base_lr
if self.backbone_pretrained:
backbone_lr *= self.lr_mult
return self.model.get_params({
'backbone': backbone_lr,
'task': base_lr
})
def save_ckpt(self, stepvalues, epoch, iters, lr, logger):
model_name = 'ckpt/model_{}.pkl'.format(int(iters))
msg = 'Snapshotting to {}'.format(model_name)
logger.info(msg)
model = {
'model': self.model.state_dict(),
'epoch': epoch,
'iters': iters,
'lr': lr,
'stepvalues': stepvalues,
'im_w': self.im_w,
'im_h': self.im_h
}
torch.save(model, model_name)
def restore_from_ckpt(self, ckpt_path, logger):
ckpt = torch.load(ckpt_path)
epoch = ckpt['epoch']
iters = ckpt['iters']
lr = ckpt['lr']
stepvalues = ckpt['stepvalues']
msg='Restoring from {}\nStart epoch: {}, total iters: {}, current lr: {}'.format(ckpt_path,\
epoch, iters, lr)
logger.info(msg)
logger.info('Model update successfully')
return stepvalues, epoch, iters, lr
def restore_from_epoch(self, start_epoch, stepvalues, lr, iters_per_epoch):
epoch_iters = int(start_epoch * iters_per_epoch)
base_lr = lr
stepvalues_new = []
for i in stepvalues:
if i > start_epoch:
stepvalues_new.append(i)
else:
lr *= self.decay_ratio
print('lr drops from {} to {}'.format(base_lr, lr))
print('Restoring from epoch {}, iters {}'.format(
start_epoch, epoch_iters))
stepvalues = stepvalues_new
return stepvalues, epoch_iters, lr
def train(self, pretrained_model=None):
coco_reader = COCODataReader(self.im_w,
self.im_h,
batch_size=self.batch_size)
coco_loader = DataLoader(coco_reader,
shuffle=True,
batch_size=self.batch_size,
num_workers=2)
self.model = FasterRCNN(self.im_w, self.im_h)
num_samples = coco_reader.__len__()
epoch_iters = num_samples // self.batch_size
num_epochs = 50
lr = 0.0002
stepvalues = [30]
num_vis_anchors = 100
config_params={'lr':lr,'epoch':0,'start_epoch':0,'num_epochs':num_epochs,'epoch_iters':epoch_iters,\
'out_size':coco_reader.out_size,'K':self.K,\
'vis_anchors_dir':'./vis_anchors','display':20,
'num_vis_anchors':num_vis_anchors}
logger = logging.getLogger(__name__)
logger.info('Load {} samples'.format(num_samples))
if pretrained_model is None:
self.model.init_weights()
else:
self.model.load_weights(model_path=pretrained_model)
self.model.cuda()
num_jumped = 0
iters_per_epoch = int(num_samples / self.batch_size)
start_epoch = 0
if pretrained_model is not None:
stepvalues, start_epoch, epoch_iters, lr = self.restore_from_ckpt(
pretrained_model, logger)
else:
stepvalues, epoch_iters, lr = self.restore_from_epoch(
start_epoch, stepvalues, lr, iters_per_epoch)
if self.backbone_pretrained: #backbone lr*=0.5
params = self.get_param_groups(lr)
optimizer = optim.SGD(params, lr=lr)
for param_group in optimizer.param_groups:
print('{} has learning rate {}'.format(param_group['key'],
param_group['lr']))
else:
optimizer = optim.SGD(self.model.parameters(), lr=lr)
logger.info('Start training...')
for epoch in range(start_epoch, num_epochs):
config_params['epoch'] = epoch
train_epoch(self.model, coco_loader, optimizer, logger,
config_params)
if len(stepvalues) > 0 and (epoch + 1) == stepvalues[0]:
lr *= self.decay_ratio
config_params['lr'] = lr
msg = 'learning rate decay: %e' % lr
logger.info(msg)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if self.backbone_pretrained and param_group[
'key'] == 'backbone':
param_group['lr'] *= self.lr_mult
stepvalues.pop(0)
self.save_ckpt(stepvalues, epoch, epoch_iters, lr, logger)
msg = 'Finish training!!\nTotal jumped batch: {}'.format(num_jumped)
logger.info(msg)
def train(self, pretrained_model=None):
coco_reader = COCODataReader(self.im_w,
self.im_h,
batch_size=self.batch_size)
coco_loader = DataLoader(coco_reader,
shuffle=True,
batch_size=self.batch_size,
num_workers=2)
self.model = FasterRCNN(self.im_w, self.im_h)
num_samples = coco_reader.__len__()
epoch_iters = num_samples // self.batch_size
num_epochs = 50
lr = 0.0002
stepvalues = [30]
num_vis_anchors = 100
config_params={'lr':lr,'epoch':0,'start_epoch':0,'num_epochs':num_epochs,'epoch_iters':epoch_iters,\
'out_size':coco_reader.out_size,'K':self.K,\
'vis_anchors_dir':'./vis_anchors','display':20,
'num_vis_anchors':num_vis_anchors}
logger = logging.getLogger(__name__)
logger.info('Load {} samples'.format(num_samples))
if pretrained_model is None:
self.model.init_weights()
else:
self.model.load_weights(model_path=pretrained_model)
self.model.cuda()
num_jumped = 0
iters_per_epoch = int(num_samples / self.batch_size)
start_epoch = 0
if pretrained_model is not None:
stepvalues, start_epoch, epoch_iters, lr = self.restore_from_ckpt(
pretrained_model, logger)
else:
stepvalues, epoch_iters, lr = self.restore_from_epoch(
start_epoch, stepvalues, lr, iters_per_epoch)
if self.backbone_pretrained: #backbone lr*=0.5
params = self.get_param_groups(lr)
optimizer = optim.SGD(params, lr=lr)
for param_group in optimizer.param_groups:
print('{} has learning rate {}'.format(param_group['key'],
param_group['lr']))
else:
optimizer = optim.SGD(self.model.parameters(), lr=lr)
logger.info('Start training...')
for epoch in range(start_epoch, num_epochs):
config_params['epoch'] = epoch
train_epoch(self.model, coco_loader, optimizer, logger,
config_params)
if len(stepvalues) > 0 and (epoch + 1) == stepvalues[0]:
lr *= self.decay_ratio
config_params['lr'] = lr
msg = 'learning rate decay: %e' % lr
logger.info(msg)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if self.backbone_pretrained and param_group[
'key'] == 'backbone':
param_group['lr'] *= self.lr_mult
stepvalues.pop(0)
self.save_ckpt(stepvalues, epoch, epoch_iters, lr, logger)
msg = 'Finish training!!\nTotal jumped batch: {}'.format(num_jumped)
logger.info(msg)
def run_network():
dataset = Dataset(args)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True)
testset = TestDataset(args)
testloader = torch.utils.data.DataLoader(testset,
batch_size=1,
shuffle=False)
print('data prepared, train data: %d, test data: %d' %
(len(dataset), len(testset)))
model = FasterRCNN()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
print('model construct completed')
# ipdb.set_trace()
# input = iter(dataloader).next()
# test = model(input[0], input[1], input[3])
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# setup log data writer
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(log_dir=args.log_dir)
for epoch in range(args.epoch):
train_losses = [0] * 5
for ii, (img, bbox, label, scale) in tqdm(enumerate(dataloader)):
# img:(batch_size, 3, h, w), bbox:(batch_size, n, 4), label:(batch, n), scale:(1)
model.train()
model.zero_grad()
img, bbox, label = img.cuda().float(), bbox.cuda(), label.cuda()
# ipdb.set_trace()
rpn_cls_locs, rpn_scores, \
gt_locs, gt_labels, \
roi_cls_locs, roi_scores, \
gt_roi_loc, gt_roi_label = model(img, bbox, label, scale.numpy())
# ------------------ RPN losses -------------------#
rpn_cls_locs.cuda()
gt_labels = torch.from_numpy(gt_labels).cuda().long()
gt_locs = torch.from_numpy(gt_locs).cuda()
rpn_cls_loss = F.cross_entropy(rpn_scores.view(-1, 2),
gt_labels,
ignore_index=-1)
rpn_loc_loss = _fast_rcnn_loc_loss(rpn_cls_locs.view(-1,
4), gt_locs,
gt_labels, args.rpn_sigma)
# ------------------ ROI losses (fast rcnn loss) -------------------#
n_simple = roi_cls_locs.shape[0]
roi_cls_locs = roi_cls_locs.view((n_simple, -1, 4))
roi_loc = roi_cls_locs[np.arange(0, n_simple), gt_roi_label]
gt_roi_label = torch.from_numpy(gt_roi_label).long().cuda()
gt_roi_loc = torch.from_numpy(gt_roi_loc).cuda()
roi_cls_loss = F.cross_entropy(roi_scores, gt_roi_label.cuda())
roi_loc_loss = _fast_rcnn_loc_loss(roi_loc.contiguous(),
gt_roi_loc, gt_roi_label.data,
args.roi_sigma)
loss = rpn_loc_loss + rpn_cls_loss + roi_loc_loss + roi_cls_loss
loss.backward()
optimizer.step()
# save loss
train_losses[0] += loss.item()
train_losses[1] += rpn_loc_loss.item()
train_losses[2] += rpn_cls_loss.item()
train_losses[3] += roi_loc_loss.item()
train_losses[4] += roi_cls_loss.item()
# plot image
if (ii == 0) or ((ii + 1) % args.plot_every) == 0:
# ipdb.set_trace()
ori_img = inverse_normalize(img[0].cpu().numpy())
# plot original image
img = vis_img(ori_img, bbox[0].cpu().numpy(),
label[0].cpu().numpy())
writer.add_image('original_images',
np.array(img).transpose(2, 0, 1), ii)
# plot pred image
_bboxes, _labels, _scores = model.predict([ori_img],
visualize=True)
if not len(_bboxes[0]) == 0:
print("pred_bboxes: {}, pre_labels: {}, pre_score: {}".
format(_bboxes, _labels, _scores))
img = vis_img(ori_img, _bboxes[0], _labels[0].reshape(-1),
_scores[0].reshape(-1))
writer.add_image('pre_images',
np.array(img).transpose(2, 0, 1), ii)
# if ii == 2 : break
# print("EPOCH:[{}/{}], roi_cls_loss: {};\n roi_loc_loss: {};\n rpn_cls_loss: {};\n rpn_loc_loss: {};\n loss: {}\n" \
# .format(epoch, args.epoch, roi_cls_loss, roi_loc_loss, rpn_cls_loss, rpn_loc_loss, loss))
# testing
pred_bboxes, pred_labels, pred_scores = list(), list(), list()
gt_bboxes, gt_labels, gt_difficults = list(), list(), list()
for ii, (imgs, sizes, gt_bboxes_, gt_labels_,
gt_difficults_) in tqdm(enumerate(testloader)):
model.eval()
with torch.no_grad():
sizes = [sizes[0][0].item(), sizes[1][0].item()]
pred_bboxes_, pred_labels_, pred_scores_ = model.predict(
imgs.to(device), [sizes])
gt_bboxes += list(gt_bboxes_.numpy())
gt_labels += list(gt_labels_.numpy())
gt_difficults += list(gt_difficults_.numpy())
pred_bboxes += pred_bboxes_
pred_labels += pred_labels_
pred_scores += pred_scores_
# if ii == 500: break
test_result = eval_detection_voc(pred_bboxes,
pred_labels,
pred_scores,
gt_bboxes,
gt_labels,
gt_difficults,
use_07_metric=True)
# plot loss and mAP
# ipdb.set_trace()
loss_name = [
'total_loss', 'rpn_loc_loss', 'rpn_cls_loss', 'roi_loc_loss',
'roi_cls_loss'
]
values = {}
for i in range(len(loss_name)):
values[loss_name[i]] = train_losses[i] / len(dataset)
writer.add_scalars('loss', values, epoch)
writer.add_scalar('test_mAP', test_result['map'], epoch)
# save checkpoints
if not os.path.exists('./checkpoints'):
os.makedirs("./checkpoints")
if epoch % 100 == 99:
torch.save(model.state_dict(),
'./checkpoints/checkpoint_{}_epochs.pth'.format(epoch))
def train(**kwargs):
opt._parse(kwargs)
# 训练数据加载器
data_set = Dataset(opt)
data_loader = data_.DataLoader(data_set,
batch_size=1,
shuffle=True,
num_workers=opt.num_workers)
# 测试数据加载器
test_set = TestDataset(opt)
test_dataloader = data_.DataLoader(test_set,
batch_size=1,
num_workers=1,
shuffle=False,
pin_memory=True)
# 网络模型
faster_rcnn = FasterRCNN()
trainer = FasterRCNNTrainer(faster_rcnn).cuda()
if opt.load_path:
trainer.load(opt.load_path)
print('load pretrained model from %s' % opt.load_path)
# 训练过程
for epoch in range(opt.epoch):
# eval_result = eval(test_dataloader, faster_rcnn)
# best_map = eval_result['map']
loss_list_roi_cls = []
loss_list_roi_loc = []
loss_list_rpn_cls = []
loss_list_rpn_loc = []
for ii, (img, bbox_, label_, scale) in tqdm(enumerate(data_loader)):
scale = at.scalar(scale)
img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda()
img, bbox, label = Variable(img), Variable(bbox), Variable(label)
loss_list = trainer.train_step(img, bbox, label, scale)
loss_list_roi_cls.append(loss_list.roi_cls_loss.detach().cpu().numpy())
loss_list_roi_loc.append(loss_list.roi_loc_loss.detach().cpu().numpy())
loss_list_rpn_cls.append(loss_list.rpn_cls_loss.detach().cpu().numpy())
loss_list_rpn_loc.append(loss_list.rpn_loc_loss.detach().cpu().numpy())
print("--------------------------")
print("curr epoch: ", epoch)
print("roi_cls loss: ", np.array(loss_list_roi_cls).mean())
print("roi_loc loss: ", np.array(loss_list_roi_loc).mean())
print("rpn_cls loss: ", np.array(loss_list_rpn_cls).mean())
print("rpn_loc loss: ", np.array(loss_list_rpn_loc).mean())
print("--------------------------")
eval_result = eval(test_dataloader, faster_rcnn)
if eval_result['map'] > best_map:
best_map = eval_result['map']
best_path = trainer.save(best_map=best_map)
if epoch == 9:
trainer.load(best_path)
trainer.faster_rcnn.scale_lr(opt.lr_decay)
head = FasterRCNNHead(n_class=len(VOC_BBOX_LABELS)+1, ratios=[0.5, 1, 2], anchor_scales=[8, 16, 32], feat_stride=16,
model=opt['pretrained_model'])
tail = FasterRCNNTail(n_class=len(VOC_BBOX_LABELS)+1, ratios=[0.5, 1, 2], anchor_scales=[8, 16, 32], feat_stride=16,
roi_size=7, model=opt['pretrained_model'])
'''
This code was written for alternating training strategy; however, we couldn't explore this strategy due to the time
constraint.
'''
if torch.cuda.is_available():
print('CUDA AVAILABLE')
faster_rcnn = FasterRCNN(head, tail).cuda()
else:
print('CUDA NOT AVAILABLE')
faster_rcnn = FasterRCNN(head, tail)
dataset = Dataset(opt)
dataloader = data_.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt['num_workers'])
valset = ValDataset(opt)
val_dataloader = data_.DataLoader(valset, batch_size=1, shuffle=False, num_workers=opt['num_workers'],
pin_memory=True)
testset = TestDataset(opt)
test_dataloader = data_.DataLoader(testset, batch_size=1, shuffle=False, num_workers=opt['test_num_workers'],
pin_memory=True)