def load_model(opt, pretrained_path): seed = int(time.time()) use_cuda = True gpus = '0' torch.manual_seed(seed) if use_cuda: os.environ['CUDA_VISIBLE_DEVICES'] = gpus torch.cuda.manual_seed(seed) # Create model model = YOWO(opt) model = model.cuda() # model = nn.DataParallel(model, device_ids=None) # in multi-gpu case model.seen = 0 checkpoint = torch.load(pretrained_path) epoch = checkpoint['epoch'] fscore = checkpoint['fscore'] model.load_state_dict(checkpoint['state_dict'], strict=False) return model, epoch, fscore
anchors = loss_options['anchors'].split(',') anchors = [float(i) for i in anchors] num_anchors = int(loss_options['num']) num_classes = opt.n_classes # Test parameters conf_thresh = 0.005 nms_thresh = 0.4 eps = 1e-5 use_cuda = True kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {} # Create model model = YOWO(opt) model = model.cuda() model = nn.DataParallel(model, device_ids=None) # in multi-gpu case print(model) # Load resume path if opt.resume_path: print("===================================================================") print('loading checkpoint {}'.format(opt.resume_path)) checkpoint = torch.load(opt.resume_path) model.load_state_dict(checkpoint['state_dict']) model.eval() print("===================================================================")
best_fscore = 0 # initialize best fscore # Test parameters nms_thresh = 0.4 iou_thresh = 0.5 if not os.path.exists(backupdir): os.mkdir(backupdir) torch.manual_seed(seed) if use_cuda: os.environ['CUDA_VISIBLE_DEVICES'] = gpus torch.cuda.manual_seed(seed) # Create model model = YOWO(opt) model = model.cuda() model = nn.DataParallel(model, device_ids=gpu_ids) # in multi-gpu case model.seen = 0 logging("============================ starting =============================") print(model) logging(f"# of GPUs: {ngpus}, batch_size: {batch_size}") parameters = get_fine_tuning_parameters(model, opt) optimizer = optim.SGD(parameters, lr=learning_rate/batch_size, momentum=momentum, dampening=0, weight_decay=decay*batch_size) kwargs = {'num_workers': num_workers, 'pin_memory': True} if use_cuda else {} # Load resume path if necessary