def load_model(opts, frame_size):
cfg_file_path = opts.model_config_dir + \
"/yolov4_" + str(frame_size) + ".cfg"
model = Darknet(cfg_file_path, inference=True)
weight_file = os.path.join(
opts.weights_dir, "yolov4_{}.pth".format(frame_size))
checkpoint = torch.load(
weight_file, map_location='cuda:{}'.format(opts.gpu_id))
model.load_state_dict(checkpoint['state_dict'])
model.eval()
if not opts.no_cuda:
model.cuda(opts.gpu_id)
# Zero grad for parameters
for param in model.parameters():
param.grad = None
return model
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.weight, n_classes=cfg.classes)
# pretrained_dict = torch.load(cfg.load)
# model_dict = model.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if
# k in model_dict} # pretrained_dict只保留了model_dict中存在的键。为什么直接load它会报错。要先给model_dict更新。
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)
total_params = sum(p.numel() for p in model.parameters())
print(f'{total_params:,} total parameters.')
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print(f'{total_trainable_params:,} training parameters.')
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(
model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device,
