def test(**kwargs): opt.load_path = \ '/home/fengkai/PycharmProjects/my-faster-rcnn/checkpoints/fasterrcnn_04231732_0.6941460588341642' opt._parse( kwargs ) #将调用函数时候附加的参数用,config.py文件里面的opt._parse()进行解释,然后获取其数据存储的路径,之后放到Dataset里面! testset = TestDataset(opt) test_dataloader = data_.DataLoader( testset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, #pin_memory=True ) #pin_memory锁页内存,开启时使用显卡的内存,速度更快 faster_rcnn = FasterRCNNVGG16() print(faster_rcnn) print('model construct completed') trainer = FasterRCNNTrainer(faster_rcnn).cuda() #判断opt.load_path是否存在,如果存在,直接从opt.load_path读取预训练模型,然后将训练数据的label进行可视化操作 trainer.load(opt.load_path) print('load pretrained model from %s' % opt.load_path) eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num) print('map is: ', str(eval_result['map']))
def main(): faster_rcnn = FasterRCNNVGG16(mask=opt.mask) trainer = FasterRCNNTrainer(faster_rcnn).cuda() assert os.path.isfile( args.load_path), f"Need valid checkpoint, {args.load_path} not found" trainer.load(args.load_path) ''' Check to make sure weights are dense ''' for n, m in trainer.named_modules(): if hasattr(m, 'sparse'): m.sparse = False for n, m in trainer.named_modules(): if hasattr(m, 'weight'): if m.weight.is_sparse: print("Weights are already sparse") return print("\n\n=========SIZE BEFORE=============") try: trainer.faster_rcnn.set_pruned() except: print("No masks.") get_size(trainer) trainer.quantize(bits=args.bits, verbose=args.verbose) print("\n\n=========SIZE AFTER==============") get_size(trainer) print("Saving a maskedmodel") trainer.save(save_path=args.save_path) print("Saving a SparseDense Model") trainer.replace_with_sparsedense() sd_file = args.save_path.split("/") sd_file[-1] = "SparseDense_" + sd_file[-1] sd_file = "/".join(sd_file) trainer.save(save_path=sd_file)
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 main(**kwargs): opt._parse(kwargs) # Initialize Logger if opt.benchmark_path is None: timestr = time.strftime('%m%d%H%M') benchmark_path = f'logs/fasterrcnn_{timestr}' for k_, v_ in kwargs.items(): benchmark_path += f'_{v_}' benchmark_path += '.log' Logger(benchmark_path, logging.INFO) logger = logging.getLogger(__name__) Logger.section_break(title='Benchmark Model') logger.info(f'User Arguments\n{opt._state_dict()}') # Load dataset dataset = TestDataset(opt, split='test') dataloader = data_.DataLoader(dataset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, pin_memory=True) logger.info(f"DATASET SIZE: {len(dataloader)}") logger.info("Using Mask VGG") if opt.mask else logger.info( "Using normal VGG16") # Construct model faster_rcnn = FasterRCNNVGG16(mask=opt.mask) trainer = FasterRCNNTrainer(faster_rcnn).cuda() Logger.section_break(title='Model') logger.info(str(faster_rcnn)) # Resume from a checkpoint if opt.load_path: assert os.path.isfile(opt.load_path),\ 'Checkpoint {} does not exist.'.format(opt.load_path) trainer.load(opt.load_path) Logger.section_break('Checkpoint') logger.info("Loaded checkpoint '{}' (epoch X)".format(opt.load_path)) # Benchmark dataset fps = AverageMeter() benchmarker = {FPS: fps} result = benchmark(benchmarker, dataloader, faster_rcnn, test_num=1000) Logger.section_break('Benchmark completed') model_parameters = filter(lambda p: p.requires_grad, faster_rcnn.parameters()) params = sum([np.prod(p.size()) for p in model_parameters]) logger.info('[PARAMETERS] {params}'.format(params=params)) logger.info('[RUN TIME] {time.avg:.3f} sec/frame'.format(time=result[FPS]))
def train(**kwargs): opt._parse(kwargs) dataset = VOC2007Detect(path=opt.root_path, is_train=True) dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.num_workers) test_dataset = VOC2007Detect(path=opt.root_path, is_train=False) test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.num_workers) print('training and testing data load finished.') print('\033[1;42m') print('---------Starting model initinalize.---------') faster_rcnn = FasterRCNNVGG16() print('model construct completed.') trainer = FasterRCNNTrainer(faster_rcnn).cuda() if opt.load_model_path: trainer.load(opt.load_model_path) print('load pretrained model from %s' % opt.load_model_path) print('len(dataloader): ', len(dataloader)) print('\033[0m') print('---------Starting training stage.---------') best_map = 0 lr_ = opt.lr for epoch in range(opt.max_epoch): trainer.reset_meters() for ii, (img, bbox_, label_, scale) in enumerate(dataloader): scale = scalar(scale) print('---------------------------------------------') print('Testing --> img.shape, bbox_.shape, label_.shape, scale: ', img.shape, bbox_.shape, label_.shape, scale) img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda() losses = trainer.train_step(img, bbox, label, scale) break eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num) lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr'] if eval_result['map'] > best_map: best_map = eval_result['map'] best_path = trainer.save(best_map=best_map)
def main(): print(opt._parse_all()) dataset = Dataset(opt) dataloader = data_.DataLoader(dataset, \ batch_size=1, \ shuffle=True, \ # pin_memory=True, num_workers=opt.num_workers) valset = TestDataset(opt, split='val') val_dataloader = data_.DataLoader(valset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, \ pin_memory=True ) testset = TestDataset(opt, split='test') test_dataloader = data_.DataLoader(testset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, \ pin_memory=True ) print(f"TRAIN SET: {len(dataloader)} | VAL SET: {len(val_dataloader)} | TEST SET: {len(test_dataloader)}") print("Using Mask VGG") if opt.mask else print("Using normal VGG16") faster_rcnn = FasterRCNNVGG16(mask=opt.mask) print('model construct completed') trainer = FasterRCNNTrainer(faster_rcnn).cuda() trainer.vis.text(dataset.db.label_names, win='labels') best_map = 0 lr_ = opt.lr start_epoch = 0 if opt.load_path: assert os.path.isfile(opt.load_path), 'Checkpoint {} does not exist.'.format(opt.load_path) checkpoint = torch.load(opt.load_path)['other_info'] if opt.use_simple: start_epoch = 0 best_map = 0 else: start_epoch = checkpoint['epoch'] best_map = checkpoint['best_map'] trainer.load(opt.load_path) print("="*30+" Checkpoint "+"="*30) print("Loaded checkpoint '{}' (epoch {})".format(opt.load_path, start_epoch)) train(opt, faster_rcnn, dataloader, val_dataloader, test_dataloader, trainer, lr_, best_map, start_epoch)
def main(): """调试用""" faster_rcnn = FasterRCNNVGG16().cuda() trainer = FasterRCNNTrainer(faster_rcnn) img = np.random.randn(3, 800, 800).astype(np.float32) img = torch.from_numpy(img[None]).cuda() bbox = np.array([[10, 20, 30, 40], [20, 30, 40, 50]]).astype(np.float32) bbox = torch.from_numpy(bbox[None]).cuda() label = np.array([1, 2], dtype=np.int32) label = torch.from_numpy(label[None]).cuda() scale = 1. losses = trainer(img, bbox, label, scale) import ipdb ipdb.set_trace() print("ok")
def __init__(self): self.frames_raw = Queue(maxsize=FaceDetector.qsize) self.frames_nhw = Queue(maxsize=FaceDetector.qsize) self.frames_ready = Queue(maxsize=FaceDetector.qsize) self.extrac_rslt = Queue(maxsize=FaceDetector.qsize) self.rpn_rslt = Queue(maxsize=FaceDetector.qsize) self.roihead_rslt = Queue(maxsize=FaceDetector.qsize) faster_rcnn = FasterRCNNVGG16(n_fg_class=1) pretrain_path = "D:/FRCNN_DATA/Faster-R-CNN-Pytorch-models/usable/fasterrcnn_04142205_0.49527486293443446" state = torch.load(pretrain_path) pretrained_model = state["model"] faster_rcnn.load_state_dict(pretrained_model) img_capture.delay(self.frames_raw, self.frames_nhw)
def main(): parser = argparse.ArgumentParser() parser.add_argument("-p", "--path") parser.add_argument("-s", "--set_id") args = parser.parse_args() valset = TestDataset(opt, set_id=args.set_id, split='val') val_dataloader = data_.DataLoader(valset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, pin_memory=True) print(f"VAL SET: {len(val_dataloader)} ") print("Using Mask VGG") if opt.mask else print("Using normal VGG16") faster_rcnn = FasterRCNNVGG16(mask=opt.mask) print('model construct completed') trainer = FasterRCNNTrainer(faster_rcnn).cuda() best_map = 0 lr_ = opt.lr if args.path: assert os.path.isfile( args.path), 'Checkpoint {} does not exist.'.format(args.path) checkpoint = torch.load(args.path)['other_info'] best_map = checkpoint['best_map'] trainer.load(args.path) print("=" * 30 + " Checkpoint " + "=" * 30) print("Loaded checkpoint '{}' ".format(args.path, best_map)) eval_result = eval(val_dataloader, faster_rcnn, test_num=1000) lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr'] # log_info = 'lr:{}, loss:{},map:{},lamr:{}'.format(str(lr_), # str(trainer.get_meter_data()), # str(eval_result['map']), # str(eval_result['lamr'])) log_info = 'lr:{}, loss:{},map:{}'.format( str(lr_), str(trainer.get_meter_data()), str(eval_result['map'])) print("Evaluation Results on Validation Set: ") print(log_info) print("\n\n") else: print("No checkpoint to evaluate is specified")
def main(): dataset = Dataset(opt) dataloader = data_.DataLoader(dataset, \ batch_size=1, \ shuffle=True, \ # pin_memory=True, num_workers=opt.num_workers) testset = TestDataset(opt, split='val') test_dataloader = data_.DataLoader(testset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, \ pin_memory=True ) print(f"TRAIN SET: {len(dataloader)} | TEST SET: {len(test_dataloader)}") faster_rcnn = FasterRCNNVGG16(mask=opt.mask) print('model construct completed') trainer = FasterRCNNTrainer(faster_rcnn).cuda() best_map = 0 lr_ = opt.lr if opt.load_path: assert os.path.isfile(opt.load_path), 'Checkpoint {} does not exist.'.format(opt.load_path) checkpoint = torch.load(opt.load_path)['other_info'] trainer.load(opt.load_path) print("="*30+" Checkpoint "+"="*30) print("Loaded checkpoint '{}' (epoch {})".format(opt.load_path, 1)) #no saved epoch, put in 1 for now if args.prune_by_std: trainer.faster_rcnn.prune_by_std(args.sensitivity) else: trainer.faster_rcnn.prune_by_percentile(q=args.percentile) prune_utils.print_nonzeros(trainer.faster_rcnn) train(opt, faster_rcnn, dataloader, test_dataloader, trainer, lr_, best_map) trainer.faster_rcnn.set_pruned() trainer.save(save_path=args.save_path) else: print("Must specify load path to pretrained model")
def train(**kwargs): opt._parse( kwargs ) #将调用函数时候附加的参数用,config.py文件里面的opt._parse()进行解释,然后获取其数据存储的路径,之后放到Dataset里面! dataset = Dataset(opt) print('load data') dataloader = data_.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.num_workers) testset = TestDataset(opt) test_dataloader = data_.DataLoader( testset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, #pin_memory=True ) #pin_memory锁页内存,开启时使用显卡的内存,速度更快 faster_rcnn = FasterRCNNVGG16() print('model construct completed') trainer = FasterRCNNTrainer(faster_rcnn).cuda() #判断opt.load_path是否存在,如果存在,直接从opt.load_path读取预训练模型,然后将训练数据的label进行可视化操作 if opt.load_path: trainer.load(opt.load_path) print('load pretrained model from %s' % opt.load_path) trainer.vis.text(dataset.dataset.label_names, win='labels') best_map = 0 lr_ = opt.lr # 之后用一个for循环开始训练过程,而训练迭代的次数opt.epoch=14也在config.py文件中都预先定义好,属于超参数 for epoch in range(opt.epoch): print('epoch {}/{}'.format(epoch, opt.epoch)) trainer.reset_meters() #首先在可视化界面重设所有数据 for ii, (img, bbox_, label_, scale) in tqdm(enumerate(dataloader)): scale = array_tool.scalar(scale) img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda() trainer.train_step(img, bbox, label, scale) if (ii + 1) % opt.plot_every == 0: if os.path.exists(opt.debug_file): ipdb.set_trace() #可视化画出loss trainer.vis.plot_many(trainer.get_meter_data()) #可视化画出groudtruth bboxes ori_img_ = inverse_normalize(array_tool.tonumpy(img[0])) gt_img = visdom_bbox(ori_img_, array_tool.tonumpy(bbox_[0]), array_tool.tonumpy(label_[0])) trainer.vis.img('gt_img', gt_img) #可视化画出预测bboxes # 调用faster_rcnn的predict函数进行预测,预测的结果保留在以_下划线开头的对象里面 _bboxes, _labels, _scores = trainer.faster_rcnn.predict( [ori_img_], visualize=True) pred_img = visdom_bbox( ori_img_, array_tool.tonumpy(_bboxes[0]), array_tool.tonumpy(_labels[0]).reshape(-1), array_tool.tonumpy(_scores[0])) trainer.vis.img('pred_img', pred_img) # 调用 trainer.vis.text将rpn_cm也就是RPN网络的混淆矩阵在可视化工具中显示出来 trainer.vis.text(str(trainer.rpn_cm.value().tolist()), win='rpn_cm') #将roi_cm也就是roihead网络的混淆矩阵在可视化工具中显示出来 trainer.vis.img( 'roi_cm', array_tool.totensor(trainer.roi_cm.conf, False).float()) eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num) trainer.vis.plot('test_map', eval_result['map']) lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr'] log_info = 'lr:{}, map:{}, loss{}'.format( str(lr_), str(eval_result['map']), str(trainer.get_meter_data())) trainer.vis.log(log_info) #将学习率以及map等信息及时显示更新 if eval_result['map'] > best_map: best_map = eval_result['map'] best_path = trainer.save(best_map=best_map) if epoch == 9: #if判断语句如果学习的epoch达到了9就将学习率*0.1变成原来的十分之一 trainer.load(best_path) trainer.faster_rcnn.scale_lr(opt.lr_decay) lr_ = lr_ * opt.lr_decay if epoch == 13: break
import os import torch as t from utils.config import opt from model.faster_rcnn_vgg16 import FasterRCNNVGG16 from trainer import FasterRCNNTrainer from data.util import read_image from utils.vis_tool import vis_bbox from utils import array_tool as at from matplotlib import pyplot as plt img = read_image('/home/fengkai/dog.jpg') img = t.from_numpy(img)[None] faster_rcnn = FasterRCNNVGG16() trainer = FasterRCNNTrainer(faster_rcnn).cuda() trainer.load('/home/fengkai/PycharmProjects/my-faster-rcnn/checkpoints/fasterrcnn_04231732_0.6941460588341642') opt.caffe_pretrain=False # this model was trained from torchvision-pretrained model _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)) plt.show()
def train(**kwargs): opt._parse(kwargs) dataset = Dataset(opt) print('load data') dataloader = data_.DataLoader(dataset, \ batch_size=1, \ shuffle=True, \ # pin_memory=True, num_workers=opt.num_workers) testset = TestDataset(opt) test_dataloader = data_.DataLoader(testset, batch_size=1, num_workers=opt.test_num_workers, shuffle=False, \ pin_memory=True ) faster_rcnn = FasterRCNNVGG16() print('model construct completed') trainer = FasterRCNNTrainer(faster_rcnn).cuda() if opt.load_path: trainer.load(opt.load_path) print('load pretrained model from %s' % opt.load_path) trainer.vis.text(dataset.db.label_names, win='labels') best_map = 0 lr_ = opt.lr for epoch in range(opt.epoch): trainer.reset_meters() for ii, (img, bbox_, label_, depth_, y_rot_, scale) in tqdm(enumerate(dataloader)): scale = at.scalar(scale) img, bbox, label, depth, y_rot = img.cuda().float(), bbox_.cuda( ), label_.cuda(), depth_.cuda(), y_rot_.cuda() trainer.train_step(img, bbox, label, depth, y_rot, scale) if (ii + 1) % opt.plot_every == 0: if os.path.exists(opt.debug_file): ipdb.set_trace() # plot loss trainer.vis.plot_many(trainer.get_meter_data()) # plot groud truth bboxes ori_img_ = inverse_normalize(at.tonumpy(img[0])) gt_img = visdom_bbox(ori_img_, at.tonumpy(bbox_[0]), at.tonumpy(label_[0])) trainer.vis.img('gt_img', gt_img) # plot predicti bboxes _bboxes, _labels, _scores, pred_depth_, pred_y_rot_ = trainer.faster_rcnn.predict( [ori_img_], visualize=True) pred_img = visdom_bbox(ori_img_, at.tonumpy(_bboxes[0]), at.tonumpy(_labels[0]).reshape(-1), at.tonumpy(_scores[0])) trainer.vis.img('pred_img', pred_img) # rpn confusion matrix(meter) trainer.vis.text(str(trainer.rpn_cm.value().tolist()), win='rpn_cm') # roi confusion matrix trainer.vis.img( 'roi_cm', at.totensor(trainer.roi_cm.conf, False).float()) eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num) print('TTTTTTTTTTTTTest_map', eval_result['map']) trainer.vis.plot('test_map', eval_result['map']) lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr'] log_info = 'lr:{}, map:{},loss:{}'.format( str(lr_), str(eval_result['map']), str(trainer.get_meter_data())) trainer.vis.log(log_info) 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) lr_ = lr_ * opt.lr_decay if epoch == 13: break
def train(**kwargs): opt._parse(kwargs) data_set = TrainDataset() print('load data.') data_loader = data_.DataLoader(data_set, batch_size=1, shuffle=True) testset = TestDataset() test_dataloader = data_.DataLoader(testset, batch_size=1, shuffle=False, pin_memory=True) faster_rcnn = FasterRCNNVGG16() print('model construct.') trainer = FasterRCNNTrainer(faster_rcnn).cuda() if opt.load_path: trainer.load(opt.load_path) print('load pretrained model from %s' % opt.load_path) lr = opt.lr best_map = 0 for epoch in range(opt.epoch): trainer.reset_meters() for ii, (img, bbox, label, scale) in tqdm(enumerate(data_loader)): img = img.cuda() trainer.train_step(img, bbox, label, scale) if (ii + 1) % opt.plot_every == 0: ipdb.set_trace() """plot loss""" trainer.vis.plot_many(trainer.get_meter_data()) """plot gt_bbox""" ori_img = inverse_normalize(img[0].cpu().numpy()) gt_img = visdom_bbox(ori_img, bbox[0].numpy(), label[0].numpy()) trainer.vis.img('gt_img', gt_img) """plot predicted bbox""" pred_bbox, pred_label, pred_score = trainer.faster_rcnn.predict( [ori_img], visualize=True) pred_img = visdom_bbox(ori_img, pred_bbox[0], pred_label[0], pred_score[0]) trainer.vis.img('pred_img', pred_img) # rpn confusion matrix(meter) trainer.vis.text(str(trainer.rpn_cm.value().tolist()), win='rpn_cm') # roi confusion matrix trainer.vis.img('roi_cm', trainer.roi_cm.conf.float().cpu()) eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num) 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) lr = lr * opt.lr_decay trainer.vis.plot('test_map', eval_result['map']) log_info = 'lr:{}, map:{},loss:{}'.format( str(lr), str(eval_result['map']), str(trainer.get_meter_data())) trainer.vis.log(log_info) if epoch == 13: print('finish!') break
from data.dataset import TrainDataset from torch.utils import data as data_ from tqdm import tqdm from model.faster_rcnn_vgg16 import decom_vgg16, VGG16RoIHead import torch from model.region_proposal_network import RegionProposalNetwork import numpy as np from model.faster_rcnn_vgg16 import FasterRCNNVGG16 from trainer import FasterRCNNTrainer data_set = TrainDataset() #data_loader = data_.DataLoader(data_set, batch_size=1, shuffle=False) img, bbox, label, scale = data_set.__getitem__(0) model = FasterRCNNVGG16().cuda() trainer = FasterRCNNTrainer(model) loss = trainer.forward( torch.from_numpy(img[None, :]).cuda(), bbox, label, scale) print(loss) """ roi_locs, roi_scores, rpn_locs, rpn_scores = model.forward(torch.from_numpy(img[None, :]).cuda()) print(roi_locs.shape) print(roi_scores.shape) print(rpn_locs.shape) print(rpn_scores.shape) """ """ extractor, classifier = decom_vgg16() feature_map = extractor.cuda()(torch.from_numpy(img[None, :]).cuda()) print(img.shape)#3, 600, 800 print(feature_map.shape)#1, 512, 37, 50
from model.faster_rcnn_vgg16 import FasterRCNNVGG16 from data.dataset import VOCBboxDataset from torch.utils.data import DataLoader from config import opt model = FasterRCNNVGG16(opt) import numpy as np import pickle global_step = 0 record_step = 10 ls = np.zeros((5)) ls_record = {} for epoch in range(opt.epoch): train_dataset = VOCBboxDataset(opt) train_num = len(train_dataset) train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True) for i,(original_img, original_bbox, img, bbox, label, scale, flip) in enumerate(train_dataloader): losses = model.train_step(img,bbox,label,scale,epoch) print('Epoch{} [{}/{}] \tTotal Loss: {:.6f}'.format(epoch, i,train_num,losses.total_loss.item())) # here can be delete global_step += 1 ls[0] += losses.rpn_loc_loss.item() ls[1] += losses.rpn_cls_loss.item()
def load_net(self): from model.faster_rcnn_vgg16 import FasterRCNNVGG16 self.net = FasterRCNNVGG16() self.net = self.net.to('cpu') self.net.load_state_dict(self.tp.get_meta('net_weight')) self.net = self.net.to(self.device)