def main():

    # Init random seeds.
    random.seed(config.RANDOM_SEED)
    np.random.seed(config.RANDOM_SEED)
    torch.manual_seed(config.RANDOM_SEED)
    torch.cuda.manual_seed(config.RANDOM_SEED)

    # Setup Tensorboard.
    print('\nsaving run in .. {}'.format(config.TRAINED_MODELS_DIR))
    if not os.path.exists(config.TRAINED_MODELS_DIR):
        os.makedirs(config.TRAINED_MODELS_DIR)
    writer = SummaryWriter(f'{config.TRAINED_MODELS_DIR}')

    # Load the Model.
    print()
    model = affnet.ResNetAffNet(pretrained=config.IS_PRETRAINED,
                                num_classes=config.NUM_CLASSES)
    model.to(config.DEVICE)

    # Load the dataset.
    train_loader, val_loader, test_loader = umd_dataset_loaders.load_umd_train_datasets(
    )

    # construct optimizer.
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params,
                                lr=config.LEARNING_RATE,
                                weight_decay=config.WEIGHT_DECAY,
                                momentum=config.MOMENTUM)
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=config.MILESTONES, gamma=config.GAMMA)

    # Main training loop.
    num_epochs = config.NUM_EPOCHS
    best_Fwb = -np.inf

    for epoch in range(0, num_epochs):
        print()

        if epoch < config.EPOCH_TO_TRAIN_FULL_DATASET:
            is_subsample = True
        else:
            is_subsample = False

        # train & val for one epoch
        model, optimizer = train_utils.train_one_epoch(
            model,
            optimizer,
            train_loader,
            config.DEVICE,
            epoch,
            writer,
            is_subsample=is_subsample)
        model, optimizer = train_utils.val_one_epoch(model,
                                                     optimizer,
                                                     val_loader,
                                                     config.DEVICE,
                                                     epoch,
                                                     writer,
                                                     is_subsample=is_subsample)
        # update learning rate.
        lr_scheduler.step()

        # # eval Fwb
        # model, Fwb = eval_utils.affnet_eval_umd(model, test_loader)
        # writer.add_scalar('eval/Fwb', Fwb, int(epoch))
        # # save best model.
        # if Fwb > best_Fwb:
        #     best_Fwb = Fwb
        #     writer.add_scalar('eval/Best_Fwb', best_Fwb, int(epoch))
        #     checkpoint_path = config.BEST_MODEL_SAVE_PATH
        #     train_utils.save_checkpoint(model, optimizer, epoch, checkpoint_path)
        #     print("Saving best model .. best Fwb={:.5f} ..".format(best_Fwb))
        #
        # checkpoint_path
        checkpoint_path = config.MODEL_SAVE_PATH + 'affnet_epoch_' + np.str(
            epoch) + '.pth'
        train_utils.save_checkpoint(model, optimizer, epoch, checkpoint_path)
def main():

    # Init random seeds.
    random.seed(config.RANDOM_SEED)
    np.random.seed(config.RANDOM_SEED)
    torch.manual_seed(config.RANDOM_SEED)
    torch.cuda.manual_seed(config.RANDOM_SEED)

    # Setup Tensorboard.
    print('\nsaving run in .. {}'.format(config.TRAINED_MODELS_DIR))
    if not os.path.exists(config.TRAINED_MODELS_DIR):
        os.makedirs(config.TRAINED_MODELS_DIR)
    writer = SummaryWriter(f'{config.TRAINED_MODELS_DIR}')

    # Load the Model.
    print()
    model = maskrcnn.ResNetMaskRCNN(pretrained=config.IS_PRETRAINED,
                                    num_classes=config.COCO_NUM_CLASSES)
    model.to(config.DEVICE)

    # Load the dataset.
    train_loader, val_loader = coco_dataset_loaders.load_coco_train_datasets()

    # Construct an optimizer.
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params,
                                lr=config.LEARNING_RATE,
                                weight_decay=config.WEIGHT_DECAY,
                                momentum=config.MOMENTUM)
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=config.MILESTONES, gamma=config.GAMMA)

    # Main training loop.
    num_epochs = config.NUM_EPOCHS
    for epoch in range(0, num_epochs):
        print()

        if epoch < config.EPOCH_TO_TRAIN_FULL_DATASET:
            is_subsample = True
        else:
            is_subsample = False

        # train & val for one epoch
        model, optimizer = train_utils.train_one_epoch(
            model,
            optimizer,
            train_loader,
            config.DEVICE,
            epoch,
            writer,
            is_subsample=is_subsample)
        model, optimizer = train_utils.val_one_epoch(model,
                                                     optimizer,
                                                     val_loader,
                                                     config.DEVICE,
                                                     epoch,
                                                     writer,
                                                     is_subsample=is_subsample)
        # update learning rate.
        lr_scheduler.step()

        # checkpoint_path
        CHECKPOINT_PATH = config.MODEL_SAVE_PATH + 'maskrcnn_epoch_' + np.str(
            epoch) + '.pth'
        train_utils.save_checkpoint(model, optimizer, epoch, CHECKPOINT_PATH)
        print(f'saved model to {CHECKPOINT_PATH} ..')
def main():

    # Init random seeds.
    random.seed(config.RANDOM_SEED)
    np.random.seed(config.RANDOM_SEED)
    torch.manual_seed(config.RANDOM_SEED)
    torch.cuda.manual_seed(config.RANDOM_SEED)

    # Setup Tensorboard.
    print('\nsaving run in .. {}'.format(config.TRAINED_MODELS_DIR))
    if not os.path.exists(config.TRAINED_MODELS_DIR):
        os.makedirs(config.TRAINED_MODELS_DIR)
    writer = SummaryWriter(f'{config.TRAINED_MODELS_DIR}')

    # Load the Model.
    print()
    # Compare Pytorch-Simple-MaskRCNN. with Torchvision MaskRCNN.
    model = model_utils.get_model_instance_segmentation(
        pretrained=config.IS_PRETRAINED, num_classes=config.NUM_CLASSES)
    model.to(config.DEVICE)

    # Load the dataset.
    train_loader, val_loader, test_loader = arl_affpose_dataset_loaders.load_arl_affpose_train_datasets(
    )

    # Construct an optimizer.
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params,
                                lr=config.LEARNING_RATE,
                                weight_decay=config.WEIGHT_DECAY,
                                momentum=config.MOMENTUM)
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=config.MILESTONES, gamma=config.GAMMA)

    # Main training loop.
    num_epochs = config.NUM_EPOCHS
    best_Fwb, best_mAP = -np.inf, -np.inf

    for epoch in range(0, num_epochs):
        print()

        # train & val for one epoch
        model, optimizer = train_utils.train_one_epoch(model,
                                                       optimizer,
                                                       train_loader,
                                                       config.DEVICE,
                                                       epoch,
                                                       writer,
                                                       is_subsample=True)
        model, optimizer = train_utils.val_one_epoch(model,
                                                     optimizer,
                                                     val_loader,
                                                     config.DEVICE,
                                                     epoch,
                                                     writer,
                                                     is_subsample=True)
        # update learning rate.
        lr_scheduler.step()

        # checkpoint_path
        checkpoint_path = config.MODEL_SAVE_PATH + 'maskrcnn_epoch_' + np.str(
            epoch) + '.pth'
        train_utils.save_checkpoint(model, optimizer, epoch, checkpoint_path)
def main():

    # Init random seeds.
    random.seed(config.RANDOM_SEED)
    np.random.seed(config.RANDOM_SEED)
    torch.manual_seed(config.RANDOM_SEED)
    torch.cuda.manual_seed(config.RANDOM_SEED)

    # Setup Tensorboard.
    print('\nsaving run in .. {}'.format(config.TRAINED_MODELS_DIR))
    if not os.path.exists(config.TRAINED_MODELS_DIR):
        os.makedirs(config.TRAINED_MODELS_DIR)
    writer = SummaryWriter(f'{config.TRAINED_MODELS_DIR}')

    # Load the Model.
    print()
    model = affnet.ResNetAffNet(pretrained=config.IS_PRETRAINED,
                                num_classes=config.NUM_CLASSES)
    model.to(config.DEVICE)
    torch.cuda.empty_cache()

    # # TODO: Freeze the backbone.
    # model = model_utils.freeze_backbone(model, verbose=True)

    # TODO: Load saved weights.
    print(
        f"\nrestoring pre-trained AffNet weights: {config.RESTORE_SYN_ARL_AFFNET_WEIGHTS} .. "
    )
    checkpoint = torch.load(config.RESTORE_SYN_ARL_AFFNET_WEIGHTS,
                            map_location=config.DEVICE)
    model.load_state_dict(checkpoint["model"])
    model.to(config.DEVICE)

    # Load the dataset.
    train_loader, val_loader, test_loader = arl_affpose_dataset_loaders.load_arl_affpose_train_datasets(
    )

    # Construct an optimizer.
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params,
                                lr=config.LEARNING_RATE,
                                weight_decay=config.WEIGHT_DECAY,
                                momentum=config.MOMENTUM)
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=config.MILESTONES, gamma=config.GAMMA)
    # # TODO: Load saved weights.
    # optimizer.load_state_dict(checkpoint["optimizer"])

    # Main training loop.
    num_epochs = config.NUM_EPOCHS
    best_Fwb, best_mAP = -np.inf, -np.inf

    for epoch in range(0, num_epochs):
        print()

        if epoch < config.EPOCH_TO_TRAIN_FULL_DATASET:
            is_subsample = True
        else:
            is_subsample = False

        # train & val for one epoch
        model, optimizer = train_utils.train_one_epoch(
            model,
            optimizer,
            train_loader,
            config.DEVICE,
            epoch,
            writer,
            is_subsample=is_subsample)
        model, optimizer = train_utils.val_one_epoch(model,
                                                     optimizer,
                                                     val_loader,
                                                     config.DEVICE,
                                                     epoch,
                                                     writer,
                                                     is_subsample=is_subsample)
        # update learning rate.
        lr_scheduler.step()

        model, mAP, Fwb = eval_utils.affnet_eval_arl_affpose(
            model, test_loader)
        # eval FwB
        writer.add_scalar('eval/Fwb', Fwb, int(epoch))
        if Fwb > best_Fwb:
            best_Fwb = Fwb
            writer.add_scalar('eval/Best_Fwb', best_Fwb, int(epoch))
            checkpoint_path = config.BEST_MODEL_SAVE_PATH
            train_utils.save_checkpoint(model, optimizer, epoch,
                                        checkpoint_path)
            print("Saving best model .. best Fwb={:.5f} ..".format(best_Fwb))
        # eval mAP
        writer.add_scalar('eval/mAP', mAP, int(epoch))
        if mAP > best_mAP:
            best_mAP = mAP
            writer.add_scalar('eval/Best_mAP', best_mAP, int(epoch))

        # checkpoint_path
        checkpoint_path = config.MODEL_SAVE_PATH + 'affnet_epoch_' + np.str(
            epoch) + '.pth'
        train_utils.save_checkpoint(model, optimizer, epoch, checkpoint_path)