def main(args): device = torch.device("cuda" if torch.cuda.is_available() and args.useCuda else "cpu") ds = pmr.datasets(args.dataset, args.dataDir, "val", train=True) indices = torch.randperm(len(ds)).tolist() d = torch.utils.data.Subset(ds, indices) model = pmr.maskrcnn_resnet50(True, len(ds.classes) + 1).to(device) model.eval() model.head.score_thresh = 0.3 if args.chptPath: checkpoint = torch.load(args.chptPath, map_location=device) model.load_state_dict(checkpoint) for p in model.parameters(): p.requires_grad_(False) iters = 3 for i, (image, target) in enumerate(d): image = image.to(device) target = {k: v.to(device) for k, v in target.items()} with torch.no_grad(): result = model(image) plt.figure(figsize=(12, 15)) pmr.show(image, result, ds.classes) if i >= iters - 1: break
def main(args): device = torch.device( "cuda" if torch.cuda.is_available() and args.use_cuda else "cpu") if device.type == "cuda": pmr.get_gpu_prop(show=True) print("\ndevice: {}".format(device)) d_test = pmr.datasets(args.dataset, args.data_dir, "val2017", train=True) # set train=True for eval print(args) num_classes = len(d_test.classes) + 1 model = pmr.maskrcnn_resnet50(False, num_classes).to(device) checkpoint = torch.load(args.ckpt_path, map_location=device) model.load_state_dict(checkpoint["model"]) print(checkpoint["eval_info"]) del checkpoint torch.cuda.empty_cache() print("evaluating only...") B = time.time() eval_output, iter_eval = pmr.evaluate(model, d_test, device, args) B = time.time() - B print(eval_output) print( "\ntotal time of this evaluation: {:.2f} s, speed: {:.2f} FPS".format( B, args.batch_size / iter_eval))
def main(args): device = torch.device( "cuda" if torch.cuda.is_available() and args.useCuda else "cpu") dataset_train = pmr.datasets(args.dataset, args.dataDir, "train", train=True) indices = torch.randperm(len(dataset_train)).tolist() d_train = torch.utils.data.Subset(dataset_train, indices) d_test = pmr.datasets(args.dataset, args.dataDir, "val", train=True) # set train=True for eval print(args) num_classes = len( d_train.dataset.classes) + 1 # including background class model = pmr.maskrcnn_resnet50(True, num_classes).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) if not args.chptPath is None: checkpoint = torch.load(args.chptPath, map_location=device) model.load_state_dict(checkpoint) start_epoch = 0 since = time.time() print("\nalready trained: {} epochs; to {} epochs".format( start_epoch, args.epochs)) for epoch in range(start_epoch, args.epochs): print("\nepoch: {}".format(epoch + 1)) A = time.time() pmr.train_one_epoch(model, optimizer, d_train, device, epoch, args) A = time.time() - A torch.save(model.state_dict(), "result/best.pth")
def main(args): device = torch.device( "cuda" if torch.cuda.is_available() and args.use_cuda else "cpu") if device.type == "cuda": pmr.get_gpu_prop(show=True) print("\ndevice: {}".format(device)) # ---------------------- prepare data loader ------------------------------- # dataset_train = pmr.datasets(args.dataset, args.data_dir, "train2017", train=True) indices = torch.randperm(len(dataset_train)).tolist() d_train = torch.utils.data.Subset(dataset_train, indices) d_test = pmr.datasets(args.dataset, args.data_dir, "val2017", train=True) # set train=True for eval args.warmup_iters = max(1000, len(d_train)) # -------------------------------------------------------------------------- # print(args) num_classes = len( d_train.dataset.classes) + 1 # including background class model = pmr.maskrcnn_resnet50(True, num_classes).to(device) 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_lambda = lambda x: 0.1**bisect.bisect([22, 26], x) start_epoch = 0 prefix, ext = os.path.splitext(args.ckpt_path) ckpts = glob.glob(prefix + "-*" + ext) ckpts.sort(key=lambda x: int( re.search(r"-(\d+){}".format(ext), os.path.split(x)[1]).group(1))) if ckpts: checkpoint = torch.load(ckpts[-1], map_location=device) # load last checkpoint model.load_state_dict(checkpoint["model"]) optimizer.load_state_dict(checkpoint["optimizer"]) start_epoch = checkpoint["epochs"] del checkpoint torch.cuda.empty_cache() since = time.time() print("\nalready trained: {} epochs; to {} epochs".format( start_epoch, args.epochs)) # ------------------------------- train ------------------------------------ # for epoch in range(start_epoch, args.epochs): print("\nepoch: {}".format(epoch + 1)) A = time.time() args.lr_epoch = lr_lambda(epoch) * args.lr print("lr_epoch: {:.4f}, factor: {:.4f}".format( args.lr_epoch, lr_lambda(epoch))) iter_train = pmr.train_one_epoch(model, optimizer, d_train, device, epoch, args) A = time.time() - A B = time.time() eval_output, iter_eval = pmr.evaluate(model, d_test, device, args) B = time.time() - B trained_epoch = epoch + 1 print("training: {:.2f} s, evaluation: {:.2f} s".format(A, B)) pmr.collect_gpu_info("maskrcnn", [1 / iter_train, 1 / iter_eval]) print(eval_output.get_AP()) pmr.save_ckpt(model, optimizer, trained_epoch, args.ckpt_path, eval_info=str(eval_output)) # it will create many checkpoint files during training, so delete some. prefix, ext = os.path.splitext(args.ckpt_path) ckpts = glob.glob(prefix + "-*" + ext) ckpts.sort(key=lambda x: int( re.search(r"-(\d+){}".format(ext), os.path.split(x)[1]).group(1))) n = 5 if len(ckpts) > n: for i in range(len(ckpts) - n): os.system("rm {}".format(ckpts[i])) # -------------------------------------------------------------------------- # print("\ntotal time of this training: {:.2f} s".format(time.time() - since)) if start_epoch < args.epochs: print("already trained: {} epochs\n".format(trained_epoch))
data_dir = "/mingback/students/jincheng/data/VOC2012/VOCdevkit/VOC2012" device = torch.device( "cuda" if torch.cuda.is_available() and use_cuda else "cpu") if device.type == "cuda": pmr.get_gpu_prop(show=True) print("\ndevice: {}".format(device)) ds = pmr.datasets(dataset, data_dir, "val2017" if dataset == "coco" else "val", train=True) indices = torch.randperm(len(ds)).tolist() d = torch.utils.data.Subset(ds, indices) model = pmr.maskrcnn_resnet50(True, len(ds.classes) + 1).to(device) model.eval() model.head.score_thresh = 0.3 if ckpt_path: checkpoint = torch.load(ckpt_path, map_location=device) model.load_state_dict(checkpoint["model"]) print(checkpoint["eval_info"]) del checkpoint torch.cuda.empty_cache() for p in model.parameters(): p.requires_grad_(False) # %% iters = 3
def main(args, config): log_dir = "./logs" log_path = os.path.join( log_dir, time.strftime("%Y-%m-%d-%H%M.log", time.localtime(time.time())) ) logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", handlers=[logging.FileHandler(log_path), logging.StreamHandler()], ) logger = logging.getLogger() device = torch.device( "cuda" if torch.cuda.is_available() and args.use_cuda else "cpu" ) if device.type == "cuda": pmr.get_gpu_prop(show=True) logger.info("\ndevice: {}".format(device)) # ---------------------- prepare data loader ------------------------------- # dataset_train = pmr.datasets( config["dataset"], config["data_dir"], "train2017" if config["dataset"] == "coco" else "train", train=True, ) print("classes = ", dataset_train.classes) indices = torch.randperm(len(dataset_train)).tolist() d_train = torch.utils.data.Subset(dataset_train, indices) d_test = pmr.datasets( config["dataset"], config["data_dir"], "val2017" if config["dataset"] == "coco" else "val", train=True, ) # set train=True for eval args.warmup_iters = max(1000, len(d_train)) # -------------------------------------------------------------------------- # print(args) num_classes = len(d_train.dataset.classes) + 1 # including background class model = pmr.maskrcnn_resnet50(True, num_classes).to(device) 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_lambda = lambda x: 0.1 ** bisect.bisect([22, 26], x) logger.info(model) logger.info("-" * 50) start_epoch = 0 prefix, ext = os.path.splitext(args.ckpt_path) ckpts = glob.glob(prefix + "-*" + ext) ckpts.sort( key=lambda x: int( re.search(r"-(\d+){}".format(ext), os.path.split(x)[1]).group(1) ) ) if ckpts: checkpoint = torch.load(ckpts[-1], map_location=device) # load last checkpoint model.load_state_dict(checkpoint["model"]) optimizer.load_state_dict(checkpoint["optimizer"]) start_epoch = checkpoint["epochs"] del checkpoint torch.cuda.empty_cache() since = time.time() logger.info( "\nalready trained: {} epochs; to {} epochs".format( start_epoch, config["epochs"] ) ) # ------------------------------- train ------------------------------------ # for epoch in range(start_epoch, config["epochs"]): logger.info("epoch: {}".format(epoch + 1)) A = time.time() args.lr_epoch = lr_lambda(epoch) * args.lr logger.info( "lr_epoch: {:.4f}, factor: {:.4f}".format(args.lr_epoch, lr_lambda(epoch)) ) iter_train = pmr.train_one_epoch( model, optimizer, d_train, device, epoch, args, logger=logger ) A = time.time() - A B = time.time() eval_output, iter_eval = pmr.evaluate(model, d_test, device, args) B = time.time() - B trained_epoch = epoch + 1 print("training: {:.2f} s, evaluation: {:.2f} s".format(A, B)) pmr.collect_gpu_info("maskrcnn", [1 / iter_train, 1 / iter_eval]) if len(list(eval_output.buffer)) > 0: # print(eval_output.get_AP()) logger.info( "bbox AP: {}; mask AP: {}".format( float(eval_output.get_AP().get("bbox AP")), float(eval_output.get_AP().get("mask AP")), ) ) pmr.save_ckpt( model, optimizer, trained_epoch, args.ckpt_path, eval_info=str(eval_output) ) # it will create many checkpoint files during training, so delete some. prefix, ext = os.path.splitext(args.ckpt_path) ckpts = glob.glob(prefix + "-*" + ext) ckpts.sort( key=lambda x: int( re.search(r"-(\d+){}".format(ext), os.path.split(x)[1]).group(1) ) ) n = 5 if len(ckpts) > n: for i in range(len(ckpts) - n): os.system("rm {}".format(ckpts[i])) # -------------------------------------------------------------------------- # logger.info("\ntotal time of this training: {:.2f} s".format(time.time() - since)) if start_epoch < args.epochs: logger.info("already trained: {} epochs\n".format(trained_epoch))