cfg.width = cfg.size
cfg.height = cfg.size
cfg.w = cfg.size
cfg.h = cfg.size
print(cfg)
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.pretrained, n_classes=cfg.classes)
if cfg.load is not None:
model.load_model(cfg.load, device)
if cfg.pretrained is not None:
model.load_model(cfg.pretrained, device, pretrained=True)
#model.load_weights(cfg.pretrained)
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,
log_step=20)
if dataset == 'patch':
from cfg_patch import Cfg
checkpoint = 'checkpoints/Yolov4_modanet_128_epoch46.pth'
elif dataset == 'modanet_whole':
pass
cfg = Cfg
cfg.gpu = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
json_gt = os.path.expanduser('~/data/datasets/modanet/Annots/modanet_instances_val_new.json')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
model.load_model(checkpoint, device)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
#model.load_state_dict(torch.load(checkpoint), strict=True)
val_dataset = YoloModanetHumanDataset(cfg.anno_path, cfg, train=False)
n_val = len(val_dataset)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=True, num_workers=0,
pin_memory=True, drop_last=False, collate_fn=val_collate)
evaluate_nms(model, val_loader, cfg, device, human_patch=True, json_gt=None)
#evaluate_nms(model, val_loader, cfg, device, human_patch=True, json_gt=json_gt)
cfg.width = cfg.size
cfg.height = cfg.size
cfg.w = cfg.size
cfg.h = cfg.size
print(cfg)
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.pretrained, n_classes=cfg.classes)
if cfg.load is not None:
model.load_model(cfg.load, device)
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, )
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')
logging.info('Saved interrupt')
try:
sys.exit(0)
