def __init__(self, viz_arch=False): self.cfg = cfg # Build model print('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) with torch.no_grad(): self.priors = Variable(self.priorbox.forward()) # Print the model architecture and parameters # if viz_arch is True: # print('Model architectures:\n{}\n'.format(self.model)) # # Utilize GPUs for computation # self.use_gpu = torch.cuda.is_available() # self.half = False # if self.use_gpu: # print('Utilize GPUs for computation') # print('Number of GPU available', torch.cuda.device_count()) # self.model.cuda() # self.priors.cuda() # cudnn.benchmark = True # # self.model = torch.nn.DataParallel(self.model).module # # Utilize half precision # self.half = cfg.MODEL.HALF_PRECISION # if self.half: # self.model = self.model.half() # self.priors = self.priors.half() # Build preprocessor and detector # self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2) self.detector = Detect(cfg.POST_PROCESS, self.priors)
def __init__(self): self.cfg = cfg self.preproc = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2) # Build model print('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) self.priors = Variable(self.priorbox.forward(), volatile=True) self.detector = Detect(cfg.POST_PROCESS, self.priors) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() if self.use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) self.model.cuda() self.priors.cuda() cudnn.benchmark = True # if torch.cuda.device_count() > 1: # self.model = torch.nn.DataParallel(self.model).module # Print the model architecture and parameters print('Model architectures:\n{}\n'.format(self.model)) num_parameters = sum([l.nelement() for l in self.model.parameters()]) print('number of parameters: {}'.format(num_parameters)) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # Set the logger self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def test(): confg_file = '/home/hby/mycode/ssds.pytorch-1/experiments/cfgs/my_ssdlite_mobilenetv2_6.yml' cfg_from_file(confg_file) model, priors = create_model(cfg.MODEL) # Utilize GPUs for computation use_gpu = torch.cuda.is_available() #half = False if use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) model.cuda() # self.model = torch.nn.DataParallel(self.model).module # Utilize half precision half = cfg.MODEL.HALF_PRECISION if half: model = model.half() pthfile = r'/home/hby/mycode/ssds.pytorch-1/experiments/models/ssd_mobilenet_v2_voc_6/ssd_lite_mobilenet_v2_voc_epoch_400.pth' checkpoint = torch.load(pthfile, map_location='cuda' if use_gpu else 'cpu') model.load_state_dict(checkpoint) #data type nchw dummy_input1 = torch.randn(1, 3, 300, 300).cuda().half() # dummy_input2 = torch.randn(1, 3, 64, 64) # dummy_input3 = torch.randn(1, 3, 64, 64) input_names = [ "actual_input_1"] output_names = [ "output1" ] # torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names) torch.onnx.export(model, dummy_input1, "ssdmobilenet_half.onnx", verbose=True, input_names=input_names, output_names=output_names)
def __init__(self): self.cfg = cfg # Load data print('===> Loading data') self.train_loader = load_data( cfg.DATASET, 'train') if 'train' in cfg.PHASE else None self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None self.visualize_loader = load_data( cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None # Build model print('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) with torch.no_grad(): self.priors = torch.Tensor(self.priorbox.forward()) self.detector = Detect(cfg.POST_PROCESS, self.priors) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() self.device = torch.device("cuda:0" if self.use_gpu else "cpu") if self.use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) # self.model.cuda() # self.priors.cuda() self.model = self.model.to(device) self.priors = self.priors.to(device) cudnn.benchmark = True # if torch.cuda.device_count() > 1: # self.model = torch.nn.DataParallel(self.model).module # Print the model architecture and parameters print('Model architectures:\n{}\n'.format(self.model)) # print('Parameters and size:') # for name, param in self.model.named_parameters(): # print('{}: {}'.format(name, list(param.size()))) # print trainable scope print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE) self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER) self.exp_lr_scheduler = self.configure_lr_scheduler( self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu) # Set the logger self.writer = SummaryWriter(log_dir=cfg.LOG_DIR) self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, input_cfg=None, root_dir=None, viz_arch=False): global cfg self.cfg = cfg if input_cfg is None else input_cfg cfg = self.cfg # Build model logger.info('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) self.priors = self.priorbox.forward() # Print the model architecture and parameters if viz_arch is True: logger.info('Model architectures:\n{}\n'.format(self.model)) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() #self.half = False self.half = cfg.MODEL.HALF_PRECISION if self.half: self.model = BN_convert_float(self.model.half()) self.priors = self.priors.half() if self.use_gpu: logger.info('Utilize GPUs for computation') logger.info('Number of GPU available: {}'.format(torch.cuda.device_count())) self.model.cuda() self.priors.cuda() cudnn.benchmark = True # self.model = torch.nn.DataParallel(self.model).module # Utilize half precision # Build preprocessor and detector self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2) self.detector = Detect(cfg.POST_PROCESS, self.priors) # Load weight: if cfg.RESUME_CHECKPOINT == '': AssertionError('RESUME_CHECKPOINT can not be empty') if cfg.RESUME_CHECKPOINT == 'latest': checkpoint_file = os.path.join(cfg.EXP_DIR, 'checkpoint_list.txt') assert os.path.exists(checkpoint_file), \ 'RESUME_CHECKPOINT set to \'latest\' but {} does not exist'.format(checkpoint_file) with open(checkpoint_file, 'r') as f: last_line = f.readlines()[-1] checkpoint_file = last_line[last_line.find(':') + 2:-1] else: checkpoint_file = cfg.RESUME_CHECKPOINT if root_dir is not None: checkpoint_file = os.path.join(root_dir, checkpoint_file) logger.info('=> loading checkpoint {:s}'.format(checkpoint_file)) checkpoint = torch.load(checkpoint_file) self.model.load_state_dict(checkpoint) # test only self.model.eval()
def __init__(self): self.cfg = cfg # Load data logger.info('Loading data') self.loaders = {} if isinstance(cfg.PHASE, (str, unicode)): cfg.PHASE = [cfg.PHASE] for phase in cfg.PHASE: self.loaders[phase] = load_data(cfg.DATASET, phase) # Build model logger.info("Building model...") self.model, self.priorbox = create_model(cfg.MODEL) with torch.no_grad(): self.priors = self.priorbox.forward() self.detector = Detect(cfg.POST_PROCESS, self.priors) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() if self.use_gpu: logger.info('Utilize GPUs for computation') logger.info('Number of GPU available: {}'.format( torch.cuda.device_count())) self.model.cuda() self.priors.cuda() cudnn.benchmark = True # if torch.cuda.device_count() > 1: # self.model = torch.nn.DataParallel(self.model).module # Print the model architecture and parameters logger.info('Model architectures:\n{}\n'.format(self.model)) logger.debug('Parameters and size:') for name, param in self.model.named_parameters(): logger.debug('{}: {}'.format(name, list(param.size()))) # print trainable scope logger.debug('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE) self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER) self.exp_lr_scheduler = self.configure_lr_scheduler( self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric loss_func = FocalLoss if cfg.MATCHER.USE_FOCAL_LOSS else MultiBoxLoss self.criterion = loss_func(cfg.MATCHER, self.priors, self.use_gpu) # Set the logger self.writer = SummaryWriter(log_dir=cfg.LOG_DIR) self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoints_kept = cfg.TRAIN.CHECKPOINTS_KEPT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, viz_arch=False): self.cfg = cfg # Build model print('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) self.priors = Variable(self.priorbox.forward(), volatile=True) # Print the model architecture and parameters if viz_arch is True: print('Model architectures:\n{}\n'.format(self.model)) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() self.device = torch.device( 'gpu') if torch.cuda.is_available() else torch.device('cpu') self.half = False if self.use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) self.model.cuda() self.priors.cuda() cudnn.benchmark = True # self.model = torch.nn.DataParallel(self.model).module # Utilize half precision self.half = cfg.MODEL.HALF_PRECISION if self.half: self.model = self.model.half() self.priors = self.priors.half() # Build preprocessor and detector self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2) self.detector = Detect(cfg.POST_PROCESS, self.priors) # Load weight: if cfg.RESUME_CHECKPOINT == '': AssertionError('RESUME_CHECKPOINT can not be empty') print('=> loading checkpoint {:s}'.format(cfg.RESUME_CHECKPOINT)) # checkpoint = torch.load(cfg.RESUME_CHECKPOINT) checkpoint = torch.load(cfg.RESUME_CHECKPOINT, map_location='gpu' if self.use_gpu else 'cpu') self.model.load_state_dict(checkpoint) # test only self.model.eval()
def __init__(self): self.cfg = cfg # Load data print('===> Loading data') self.train_loader = load_data(cfg.DATASET, 'train') if 'train' in cfg.PHASE else None self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None self.visualize_loader = load_data(cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None if self.train_loader and hasattr(self.train_loader.dataset, "num_classes"): cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES=self.train_loader.dataset.num_classes elif self.eval_loader and hasattr(self.eval_loader.dataset, "num_classes"): cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES=self.eval_loader.dataset.num_classes elif self.test_loader and hasattr(self.test_loader.dataset, "num_classes"): cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES = self.test_loader.dataset.num_classes elif self.visualize_loader and hasattr(self.visualize_loader.dataset, "num_classes"): cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES = self.visualize_loader.dataset.num_classes # Build model print('===> Building model, num_classes is '+str(cfg.MODEL.NUM_CLASSES)) self.model, self.priorbox = create_model(cfg.MODEL,cfg.LOSS.CONF_DISTR) self.priors = Variable(self.priorbox.forward(), volatile=True) self.detector = Detect(cfg.POST_PROCESS, self.priors) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() if self.use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) self.model.cuda() self.priors.cuda() if torch.cuda.device_count() > 1: print('-----DataParallel-----------') self.model = torch.nn.DataParallel(self.model) self.model.cuda() #self.dp_model = torch.nn.DataParallel(self.model) #self.model = torch.nn.DataParallel(self.model).module #self.model = self.dp_model.module cudnn.benchmark = True # Print the model architecture and parameters #print('Model architectures:\n{}\n'.format(self.model)) # print('Parameters and size:') # for name, param in self.model.named_parameters(): # print('{}: {}'.format(name, list(param.size()))) # print trainable scope print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE) self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER) self.exp_lr_scheduler = self.configure_lr_scheduler(self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric #self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu) self.criterion = FocalLoss(cfg.MATCHER, self.priors, self.use_gpu, cfg.LOSS) # Set the logger self.writer = SummaryWriter(log_dir=cfg.LOG_DIR) self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.pretrained= cfg.PRETRAINED self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self): self.cfg = cfg # set up logging if not os.path.exists(cfg.LOG_DIR): os.mkdir(cfg.LOG_DIR) # utils.setup_logging(os.path.join(cfg.LOG_DIR, "log.txt")) logging.basicConfig(filename=os.path.join(cfg.LOG_DIR, "log.txt"), level=logging.INFO) # Load data logging.info('===> Loading data, phase %s' % cfg.PHASE) self.train_loader = load_data( cfg.DATASET, 'train') if 'train' in cfg.PHASE else None self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None self.visualize_loader = load_data( cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None # Build model logging.info('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) self.priors = self.priorbox.forward() self.detector = Detect(cfg.POST_PROCESS, self.priors) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() if self.use_gpu: logging.info('Utilize GPUs for computation') logging.info('Number of GPU available: %d' % torch.cuda.device_count()) self.model.cuda() self.priors.cuda() cudnn.benchmark = True # if torch.cuda.device_count() > 1: # self.model = torch.nn.DataParallel(self.model).module # Print the model architecture and parameters # logging.info('Model architectures:\n{}\n'.format(self.model)) # logging.info('Parameters and size:') # for name, param in self.model.named_parameters(): # logging.info('{}: {}'.format(name, list(param.size()))) # logging.info trainable scope logging.info('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE) self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER) self.exp_lr_scheduler = self.configure_lr_scheduler( self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu) # Set the logger self.writer = SummaryWriter(log_dir=cfg.LOG_DIR) self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, ifTrain=True): self.cfg = cfg # Load data print('===> Loading data') self.ifTrain = ifTrain if self.ifTrain: self.train_loader = load_data( cfg.DATASET, 'train') if 'train' in cfg.PHASE else None #self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None else: test_image_dir = os.path.join('./data/', 'ship_test_v2') # transforms = transform.Compose([transform.Lambda(lambda x: cv2.cvtColor(np.asarray(x),cv2.COLOR_RGB2BGR)),transform.Resize([300,300]), transform.ToTensor()]) # test_set = torchvision.datasets.ImageFolder(test_image_dir, transform = transforms) # self.test_loader = torch.utils.data.DataLoader(test_set,batch_size=8,shuffle=False,num_workers=8) self.train_loader = load_data( cfg.DATASET, 'train') if 'train' in cfg.PHASE else None #self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None self.visualize_loader = load_data( cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None # Build model print('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) self.priors = Variable(self.priorbox.forward(), volatile=True) self.detector = Detect(cfg.POST_PROCESS, self.priors) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() #self.use_gpu = False if self.use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) self.model.cuda() self.priors.cuda() cudnn.benchmark = True if torch.cuda.device_count() > 1: self.model = torch.nn.DataParallel(self.model).module #device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") #device = torch.device("cpu") #self.model = self.model.to(device) # Print the model architecture and parameters print('Model architectures:\n{}\n'.format(self.model)) # print('Parameters and size:') # for name, param in self.model.named_parameters(): # print('{}: {}'.format(name, list(param.size()))) # print trainable scope print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE) # print('trainable_param ', trainable_param) self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER) self.exp_lr_scheduler = self.configure_lr_scheduler( self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric # print('priors ', self.priors) self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu) # Set the logger self.writer = SummaryWriter(log_dir=cfg.LOG_DIR) self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self): self.cfg = cfg self.use_gpu = torch.cuda.is_available() if self.use_gpu: torch.set_default_tensor_type('torch.cuda.FloatTensor') # Load data print('===> Loading data') if 'train_mimic' == cfg.PHASE[0] or 'train' == cfg.PHASE[0]: self.train_loader_1 = load_data(cfg.DATASET, 'train') print(cfg.DATASET.DATASET, len(self.train_loader_1)) if cfg.DATASET2.DATASET in cfg.DATASET.DATASETS: self.train_loader_2 = load_data(cfg.DATASET2, 'train') print(cfg.DATASET2.DATASET, len(self.train_loader_1)) else: self.train_loader_2 = None self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE1 else None self.corr_loader = load_data( cfg.DATASET, 'correlation') if 'correlation' in cfg.PHASE1 else None print('===> Building model') self.model, self.priorbox, feature_maps = create_model(cfg.MODEL) sizes = [] boxes = [] for maps in feature_maps: sizes.append(maps[0] * maps[1]) for box in cfg.MODEL.ASPECT_RATIOS: boxes.append(len(box) * 2) self.priors = Variable(self.priorbox.forward(), volatile=True) self.detector = Detect_fast(cfg.POST_PROCESS, self.priors) self.Discriminator = Discriminator(cfg.LOG_DIR, cfg.MODEL.NETS, cfg.MODEL_MIMIC.NETS, cfg.DISCTRIMINATOR, sizes, boxes) self.Correlation = Correlation(cfg.CORRELATION, sizes, boxes) self.Trainer = Trainer() self.Traier_mimic = Traier_mimic(cfg.TRAIN_MIMIC, sizes, boxes, cfg.DISCTRIMINATOR.TYPE) self.Tester = Tester(cfg.POST_PROCESS) if 'train_mimic' == cfg.PHASE[0] or 'correlation' == cfg.PHASE[0]: self.model_mimic = create_model(cfg.MODEL_MIMIC) self.model_mimic.load_state_dict( torch.load(cfg.MODEL_MIMIC.WEIGHTS)) self.DNet = self.Discriminator.create_discriminator() # Utilize GPUs for computation if self.use_gpu: print('Utilize GPUs for computation') print('Number of GPU available', torch.cuda.device_count()) cudnn.benchmark = True # self.model_mimic = torch.nn.DataParallel(self.model_mimic) # for i in range(len(self.DNet)): # self.DNet[i] = torch.nn.DataParallel(self.DNet[i]) self.model.cuda() self.priors.cuda() if 'train_mimic' == cfg.PHASE[0] or 'correlation' == cfg.PHASE[0]: self.model_mimic.cuda() for i in range(len(self.DNet)): self.DNet[i] = self.DNet[i].cuda() # if 'train_mimic' == cfg.PHASE[0]: # print('Model mimim architectures:\n{}\n'.format(self.model_mimic)) # for i in range(len(self.DNet)): # print('Hello') # print(self.DNet[i]) # print('Parameters and size:') # for name, param in self.model.named_parameters(): # print('{}: {}'.format(name, list(param.size()))) # print trainable scope # print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) if 'train_mimic' == cfg.PHASE[0] or 'correlation' == cfg.PHASE[0]: self.optimizer = self.configure_optimizer(self.model, cfg.TRAIN.OPTIMIZER) self.DNet_optim = [] for i in range(len(self.DNet)): self.DNet_optim.append( self.configure_optimizer(self.DNet[i], cfg.DISCTRIMINATOR.OPTIMIZER)) self.optimizer_GENERATOR = self.configure_optimizer( self.model, cfg.TRAIN_MIMIC.OPTIMIZER) self.exp_lr_scheduler_g = self.configure_lr_scheduler( self.optimizer_GENERATOR, cfg.TRAIN.LR_SCHEDULER) else: self.optimizer = self.configure_optimizer(self.model, cfg.TRAIN.OPTIMIZER) self.phase = cfg.PHASE self.exp_lr_scheduler = self.configure_lr_scheduler( self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu) self.criterion_GaN = nn.BCELoss() self.pos = POSdata(cfg.MATCHER, self.priors, self.use_gpu) # Set the logger self.writer = SummaryWriter(log_dir=cfg.LOG_DIR) if not os.path.exists(cfg.LOG_DIR): os.mkdir(cfg.LOG_DIR) shutil.copyfile('./lib/utils/config_parse.py', cfg.LOG_DIR + 'hiperparameters.py') a = os.listdir(cfg.LOG_DIR) for i in range(1, 100): if not 'Correlation_' + str(i) + '.txt' in a: self.logger = cfg.LOG_DIR + 'Correlation_' + str(i) + '.txt' self.loglosses = cfg.LOG_DIR + 'Correlation_loss_' + str( i) + '.txt' break f = open(self.logger, 'w') f.close() f = open(self.loglosses, 'w') f.close() self.output_dir = cfg.LOG_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX self.model.loc.apply(self.weights_init) self.model.conf.apply(self.weights_init) self.model.extras.apply(self.weights_init) if 'train' == cfg.PHASE[0]: if torch.cuda.device_count() > 1: self.model = torch.nn.DataParallel(self.model)
def run_detection(data_dir, coco_gt, im_ids): model, priorbox = create_model(cfg.MODEL) priors = Variable(priorbox.forward(), volatile=True) detector = Detect(cfg.POST_PROCESS, priors) # Utilize GPUs for computation model.cuda() priors.cuda() cudnn.benchmark = True preprocess = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2) resume_checkpoint(model, args.weight) num_classes = 2 results = [] time_all = [] time_per_step = { "nms_time": [], "cpu_tims": [], "scores_time": [], "box_time": [], "gpunms_time": [], "base_time": [], "extra_time": [], "head_time": [] } for i, index in enumerate(im_ids): # load img print('evaluating image {}/{}'.format(i, len(im_ids))) im_data = coco_gt.loadImgs(ids=index)[0] img = cv2.imread(os.path.join(data_dir, 'frames', im_data['file_name'])) scale = [img.shape[1], img.shape[0], img.shape[1], img.shape[0]] img_shape = img.shape images = Variable(preprocess(img)[0].unsqueeze(0).cuda(), volatile=True) img_dict = {'version': 0, 'time': 0., 'camera_id': 0, 'image': img} # run detection model torch.cuda.synchronize() time_all_start = time.perf_counter() # forward out = model(images, phase='eval') # detect detections = detector.forward(out) torch.cuda.synchronize() time_all_end = time.perf_counter() time_all.append(1000 * (time_all_end - time_all_start)) scores = [] cls_boxes = [] for det in detections[0][1]: if det[0] > 0: d = det.cpu().numpy() score, box = d[0], d[1:] box *= scale scores.append(score) cls_boxes.append(box) #print(score) #print(box) output_blobs = {} output_blobs['scores'] = scores output_blobs['cls_boxes'] = cls_boxes print(np.array(cls_boxes).shape) output_dict = postprocess(output_blobs, 1., img_shape) if len(output_dict['people']) == 0: continue # save result entry_index = 0 for person in output_dict['people']: entry_result = { "image_id": index, "category_id": 1, "bbox": person['bbox_ltwh'].tolist(), "score": person['score'] } results.append(entry_result) # save results as json file with open(json_dt, 'w') as f: json.dump(results, f) print('detection results saved in {}'.format(json_dt)) print('average running time: {}ms'.format(sum(time_all) / len(time_all)))
def train_epoch(self, model, data_loader, optimizer, criterion, writer, epoch, use_gpu): model.train() epoch_size = len(data_loader) batch_iterator = iter(data_loader) loc_loss = 0 conf_loss = 0 l2_loss = 0 l2_loss_weight = 0.7 _t = Timer() #prepare for teacher model teacher_model, teacher_priorbox = create_model(cfg.TEACHER_MODEL) #teacher_priors = Variable(teacher_priorbox.forward(), volatile=True) #teacher_detector = Detect(cfg.POST_PROCESS, teacher_priors) # Utilize GPUs for computation teacher_model.cuda() #teacher_priors.cuda() cudnn.benchmark = True #teacher_preprocess = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2) weight = '/home/zhihuai/yolo/experiments/models/darknet_53_yolo_v3_aifi_three_60ep/yolo_v3_darknet_53_coco_epoch_60.pth' self.resume_checkpoint_teacher(teacher_model, weight) for iteration in iter(range((epoch_size))): images, targets = next(batch_iterator) if use_gpu: images = Variable(images.cuda()) targets = [ Variable(anno.cuda(), volatile=True) for anno in targets ] else: images = Variable(images) targets = [Variable(anno, volatile=True) for anno in targets] _t.tic() # forward out, student_out = model(images, phase='distill_student') teacher_out = teacher_model(images, phase='distill_teacher') teacher_feature_1 = teacher_out[0] student_feature_1 = student_out[0] teacher_feature_1.detach_() student_feature_1.detach_() loss_l2_1 = self.l2_loss(student_feature_1, teacher_feature_1) teacher_feature_2 = teacher_out[1] student_feature_2 = student_out[1] teacher_feature_2.detach_() student_feature_2.detach_() loss_l2_2 = self.l2_loss(student_feature_2, teacher_feature_2) # backprop optimizer.zero_grad() loss_l, loss_c = criterion(out, targets) # some bugs in coco train2017. maybe the annonation bug. if loss_l.data[0] == float("Inf"): continue loss = loss_l + loss_c + l2_loss_weight * loss_l2_1 + ( 1 - l2_loss_weight) * loss_l2_2 loss.backward() optimizer.step() time = _t.toc() loc_loss += loss_l.data[0] conf_loss += loss_c.data[0] l2_loss += l2_loss_weight * loss_l2_1.data[0] + ( 1 - l2_loss_weight) * loss_l2_2.data[0] # log per iter log = '\r==>Train: || {iters:d}/{epoch_size:d} in {time:.3f}s [{prograss}] || loc_loss: {loc_loss:.4f} cls_loss: {cls_loss:.4f} l2_loss: {l2_loss:.4f}\r'.format( prograss='#' * int(round(10 * iteration / epoch_size)) + '-' * int(round(10 * (1 - iteration / epoch_size))), iters=iteration, epoch_size=epoch_size, time=time, loc_loss=loss_l.data[0], cls_loss=loss_c.data[0], l2_loss=l2_loss_weight * loss_l2_1.data[0] + (1 - l2_loss_weight) * loss_l2_2.data[0]) sys.stdout.write(log) sys.stdout.flush() # log per epoch sys.stdout.write('\r') sys.stdout.flush() lr = optimizer.param_groups[0]['lr'] log = '\r==>Train: || Total_time: {time:.3f}s || loc_loss: {loc_loss:.4f} conf_loss: {conf_loss:.4f} l2_loss: {l2_loss:.4f} || lr: {lr:.6f}\n'.format( lr=lr, time=_t.total_time, loc_loss=loc_loss / epoch_size, conf_loss=conf_loss / epoch_size, l2_loss=l2_loss / epoch_size) sys.stdout.write(log) sys.stdout.flush() # log for tensorboard writer.add_scalar('Train/loc_loss', loc_loss / epoch_size, epoch) writer.add_scalar('Train/conf_loss', conf_loss / epoch_size, epoch) writer.add_scalar('Train/lr', lr, epoch)
def __init__(self, phase="train"): self.cfg = cfg creat_log(self.cfg, phase=phase) for k, v in cfg.items(): print(k, ": ", v) log_str = '\rEpoch {k}: {v}'.format(k=k, v=v) logging.info(log_str) # Load data print('===> Loading data') logging.info('===> Loading data') self.train_loader = load_data( cfg.DATASET, 'train') if 'train' in cfg.PHASE else None self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None self.visualize_loader = load_data( cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None # Build model print('===> Building model') logging.info('===> Building model') self.model, self.priorbox = create_model(cfg.MODEL) with torch.no_grad(): self.priors = Variable(self.priorbox.forward()) self.detector = Detect(cfg.POST_PROCESS, self.priors) os.makedirs(self.cfg['EXP_DIR'], exist_ok=True) # Utilize GPUs for computation self.use_gpu = torch.cuda.is_available() print('Model architectures:\n{}\n'.format(self.model)) logging.info('Model architectures:\n{}\n'.format(self.model)) from lib.utils.torchsummary import summary summary_text = summary( self.model.cuda(), (3, cfg.MODEL.IMAGE_SIZE[0], cfg.MODEL.IMAGE_SIZE[1])) logging.info('\n'.join(summary_text)) # num_params = 0 # for name, param in self.model.named_parameters(): # num_params += param.numel() # # print(name, param.size(), param.numel()) # print("%40s %20s %20s" % (name, num_params, param.numel())) # print(num_params/1e4) # df = torch_summarize_df(input_size=(3, 512, 512), model=self.model) # # df['name'], list(df['class_name']), df['input_shape'], df["output_shape"], list(df['nb_params']) # print(df) # for name, param in self.model.named_parameters(): # print(name, param.size()) # from thop import profile # # flops, params = profile(self.model, input_size=(1, 3, 512, 128)) # count = 0 # for p in self.model.parameters(): # count += p.data.nelement() # self.multi_gpu = True self.multi_gpu = False if self.use_gpu: print('Utilize GPUs for computation') logging.info('Utilize GPUs for computation') # print('Number of GPU available', torch.cuda.device_count()) self.model.cuda() self.priors.cuda() cudnn.benchmark = True # os.environ['CUDA_VISIBLE_DEVICES'] = "4,5,6,7" # "0,1,2,3,4,5,6,7" if torch.cuda.device_count() > 1 and self.multi_gpu: self.model = torch.nn.DataParallel(self.model.cuda()) cudnn.benchmark = True # self.model = torch.nn.DataParallel(self.model).module # Print the model architecture and parameters # print('Parameters and size:') # for name, param in self.model.named_parameters(): # print('{}: {}'.format(name, list(param.size()))) # print trainable scope print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) logging.info('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE)) trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE) self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER) self.exp_lr_scheduler = self.configure_lr_scheduler( self.optimizer, cfg.TRAIN.LR_SCHEDULER) self.max_epochs = cfg.TRAIN.MAX_EPOCHS # metric self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu) # Set the logger self.writer = SummaryWriter(logdir=cfg.LOG_DIR) self.output_dir = cfg.EXP_DIR self.checkpoint = cfg.RESUME_CHECKPOINT self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX