Пример #1
0
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
Пример #2
0
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("===================================================================")



def get_clip(root, imgpath, train_dur, dataset):