from core.region_loss import RegionLoss, RegionLoss_Ava
from core.model import YOWO, get_fine_tuning_parameters
####### Load configuration arguments
# ---------------------------------------------------------------
args = parser.parse_args()
cfg = parser.load_config(args)
####### Check backup directory, create if necessary
# ---------------------------------------------------------------
if not os.path.exists(cfg.BACKUP_DIR):
os.makedirs(cfg.BACKUP_DIR)
####### Create model
# ---------------------------------------------------------------
model = YOWO(cfg)
model = model.cuda()
model = nn.DataParallel(model, device_ids=None) # in multi-gpu case
# print(model)
pytorch_total_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
logging(
'Total number of trainable parameters: {}'.format(pytorch_total_params))
seed = int(time.time())
torch.manual_seed(seed)
use_cuda = True
if use_cuda:
os.environ[
'CUDA_VISIBLE_DEVICES'] = '0' # TODO: add to config e.g. 0,1,2,3
torch.cuda.manual_seed(seed)
from datasets.ava_eval_helper import read_labelmap
from datasets.meters import AVAMeter
from core.optimization import *
from cfg import parser
from core.utils import *
from core.region_loss import RegionLoss, RegionLoss_Ava
from core.model import YOWO, get_fine_tuning_parameters
####### Load configuration arguments
# ---------------------------------------------------------------
args = parser.parse_args()
cfg = parser.load_config(args)
####### Create model
# ---------------------------------------------------------------
model = YOWO(cfg)
model = model.cuda()
model = nn.DataParallel(model, device_ids=None) # in multi-gpu case
####### Load resume path if necessary
# ---------------------------------------------------------------
if cfg.TRAIN.RESUME_PATH:
print(
"===================================================================")
print('loading checkpoint {}'.format(cfg.TRAIN.RESUME_PATH))
checkpoint = torch.load(cfg.TRAIN.RESUME_PATH)
cfg.TRAIN.BEGIN_EPOCH = checkpoint['epoch'] + 1
best_score = checkpoint['score']
model.load_state_dict(checkpoint['state_dict'])
print("Loaded model score: ", checkpoint['score'])
print(