def load_model(cfg_file, testing=False): merge_cfg_from_file(cfg_file) assert_and_infer_cfg() model_def = import_from_file(cfg.MODEL.FILE) model_def = model_def.PointNetDet input_channels = 3 if not cfg.DATA.WITH_EXTRA_FEAT else 4 dataset_name = cfg.DATA.DATASET_NAME assert dataset_name in DATASET_INFO datset_category_info = DATASET_INFO[dataset_name] NUM_VEC = len( datset_category_info.CLASSES) # rgb category as extra feature vector NUM_CLASSES = len(cfg.MODEL.CLASSES) model = model_def(input_channels, num_vec=NUM_VEC, num_classes=NUM_CLASSES) model = model.cuda() if testing: if os.path.isfile(cfg.TEST.WEIGHTS): checkpoint = torch.load(cfg.TEST.WEIGHTS, map_location={'cuda:1': 'cuda:0'}) if 'state_dict' in checkpoint: # print(checkpoint['state_dict']) new_state_dict = OrderedDict() for k, v in checkpoint['state_dict'].items(): name = k[7:] # remove `module.` new_state_dict[name] = v model.load_state_dict(new_state_dict) logging.info("=> loaded checkpoint '{}' (epoch {})".format( cfg.TEST.WEIGHTS, checkpoint['epoch'])) else: model.load_state_dict(checkpoint) logging.info("=> loaded checkpoint '{}')".format( cfg.TEST.WEIGHTS)) else: logging.error("=> no checkpoint found at '{}'".format( cfg.TEST.WEIGHTS)) assert False return model
from configs.config import assert_and_infer_cfg # from src.frustum_convnet.utils.utils import import_from_file from pathlib import Path import sys import importlib import shutil from pyntcloud import PyntCloud # from src.calib import Calib from utils.utils import import_from_file use_cam = True merge_cfg_from_file('cfgs/det_sample_waymo.yaml') assert_and_infer_cfg() fr_weigths_path = 'pretrained_models/car_waymo/model_best.pth' fr_model_def = import_from_file('models/det_base_onnx.py') fr_model_def = fr_model_def.PointNetDet input_channels = 3 NUM_VEC = 0 NUM_CLASSES = 2 fr_model = fr_model_def(input_channels, num_vec=NUM_VEC, num_classes=NUM_CLASSES) fr_model = torch.nn.DataParallel(fr_model) if os.path.isfile(fr_weigths_path): checkpoint = torch.load(fr_weigths_path) if 'state_dict' in checkpoint: fr_model.load_state_dict(checkpoint['state_dict'])
def main(): # parse arguments args = parse_args() if args.cfg_file is not None: merge_cfg_from_file(args.cfg_file) if args.opts is not None: merge_cfg_from_list(args.opts) assert_and_infer_cfg() if not os.path.exists(cfg.OUTPUT_DIR): os.makedirs(cfg.OUTPUT_DIR) # set logger cfg_name = os.path.basename(args.cfg_file).split('.')[0] log_file = '{}_{}_train.log'.format(cfg_name, time.strftime('%Y-%m-%d-%H-%M')) log_file = os.path.join(cfg.OUTPUT_DIR, log_file) logger = get_logger(log_file) logger.info(pprint.pformat(args)) logger.info('config:\n {}'.format(pprint.pformat(cfg))) # set visualize logger logger_train = None logger_val = None if cfg.USE_TFBOARD: from utils.logger import Logger logger_dir = os.path.join(cfg.OUTPUT_DIR, 'tb_logger', 'train') if not os.path.exists(logger_dir): os.makedirs(logger_dir) logger_train = Logger(logger_dir) logger_dir = os.path.join(cfg.OUTPUT_DIR, 'tb_logger', 'val') if not os.path.exists(logger_dir): os.makedirs(logger_dir) logger_val = Logger(logger_dir) # import dataset set_random_seed() logging.info(cfg.DATA.FILE) dataset_def = import_from_file(cfg.DATA.FILE) collate_fn = dataset_def.collate_fn dataset_def = dataset_def.ProviderDataset train_dataset = dataset_def(cfg.DATA.NUM_SAMPLES, split=cfg.TRAIN.DATASET, one_hot=True, random_flip=True, random_shift=True, extend_from_det=cfg.DATA.EXTEND_FROM_DET) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=True, num_workers=cfg.NUM_WORKERS, pin_memory=True, drop_last=True, collate_fn=collate_fn) val_dataset = dataset_def(cfg.DATA.NUM_SAMPLES, split=cfg.TEST.DATASET, one_hot=True, random_flip=False, random_shift=False, extend_from_det=cfg.DATA.EXTEND_FROM_DET) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=cfg.TEST.BATCH_SIZE, shuffle=False, num_workers=cfg.NUM_WORKERS, pin_memory=True, drop_last=False, collate_fn=collate_fn) logging.info('training: sample {} / batch {} '.format( len(train_dataset), len(train_loader))) logging.info('validation: sample {} / batch {} '.format( len(val_dataset), len(val_loader))) logging.info(cfg.MODEL.FILE) model_def = import_from_file(cfg.MODEL.FILE) model_def = model_def.PointNetDet input_channels = 3 if not cfg.DATA.WITH_EXTRA_FEAT else cfg.DATA.EXTRA_FEAT_DIM # NUM_VEC = 0 if cfg.DATA.CAR_ONLY else 3 dataset_name = cfg.DATA.DATASET_NAME assert dataset_name in DATASET_INFO datset_category_info = DATASET_INFO[dataset_name] NUM_VEC = len( datset_category_info.CLASSES) # rgb category as extra feature vector NUM_CLASSES = cfg.MODEL.NUM_CLASSES model = model_def(input_channels, num_vec=NUM_VEC, num_classes=NUM_CLASSES) logging.info(pprint.pformat(model)) if cfg.NUM_GPUS > 1: model = torch.nn.DataParallel(model) model = model.cuda() parameters_size = 0 for p in model.parameters(): parameters_size += p.numel() logging.info('parameters: %d' % parameters_size) logging.info('using optimizer method {}'.format(cfg.TRAIN.OPTIMIZER)) if cfg.TRAIN.OPTIMIZER == 'adam': optimizer = optim.Adam(model.parameters(), lr=cfg.TRAIN.BASE_LR, betas=(0.9, 0.999), weight_decay=cfg.TRAIN.WEIGHT_DECAY) elif cfg.TRAIN.OPTIMIZER == 'sgd': optimizer = optim.SGD(model.parameters(), lr=cfg.TRAIN.BASE_LR, momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY) else: assert False, 'Not support now.' # miles = [math.ceil(num_epochs*3/8), math.ceil(num_epochs*6/8)] # assert isinstance(LR_SETP, list) LR_STEPS = cfg.TRAIN.LR_STEPS LR_DECAY = cfg.TRAIN.GAMMA if len(LR_STEPS) > 1: lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=LR_STEPS, gamma=LR_DECAY) else: lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=LR_STEPS[0], gamma=LR_DECAY) best_prec1 = 0 best_epoch = 0 start_epoch = 0 # optionally resume from a checkpoint if cfg.RESUME: if os.path.isfile(cfg.TRAIN.WEIGHTS): checkpoint = torch.load(cfg.TRAIN.WEIGHTS) start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] best_epoch = checkpoint['best_epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint '{}' (epoch {})".format( cfg.TRAIN.WEIGHTS, checkpoint['epoch'])) else: logger.error("=> no checkpoint found at '{}'".format( cfg.TRAIN.WEIGHTS)) # resume from other pretrained model if start_epoch == cfg.TRAIN.MAX_EPOCH: start_epoch = 0 best_prec1 = 0 best_epoch = 0 if cfg.EVAL_MODE: validate(val_loader, model, start_epoch, logger_val) return MAX_EPOCH = cfg.TRAIN.MAX_EPOCH for n in range(start_epoch, MAX_EPOCH): train(train_loader, model, optimizer, lr_scheduler, n, logger_train) ious_gt = validate(val_loader, model, n, logger_val) prec1 = ious_gt is_best = False if prec1 > best_prec1: best_prec1 = prec1 best_epoch = n + 1 is_best = True logging.info( 'Best model {:04d}, Validation Accuracy {:.6f}'.format( best_epoch, best_prec1)) save_data = { 'epoch': n + 1, 'state_dict': model.state_dict() if cfg.NUM_GPUS == 1 else model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'best_prec1': best_prec1, 'best_epoch': best_epoch } if (n + 1) % 5 == 0 or (n + 1) == MAX_EPOCH: torch.save( save_data, os.path.join(cfg.OUTPUT_DIR, 'model_%04d.pth' % (n + 1))) if is_best: torch.save(save_data, os.path.join(cfg.OUTPUT_DIR, 'model_best.pth')) if (n + 1) == MAX_EPOCH: torch.save(save_data, os.path.join(cfg.OUTPUT_DIR, 'model_final.pth')) logging.info('Best model {:04d}, Validation Accuracy {:.6f}'.format( best_epoch, best_prec1))