params.py

from utils.logger import setup_logger
from data import make_dataloader
from modeling import make_model
import random
import torch
import numpy as np
import os
import argparse
from config import cfg


def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True



if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="EDITOR Training")
    parser.add_argument(
        "--config_file", default="", help="path to config file", type=str
    )

    parser.add_argument("opts", help="Modify config options using the command-line", default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument("--local_rank", default=0, type=int)
    args = parser.parse_args()
    if args.config_file != "":
        cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.MODEL.CLIP_FROZEN = False
    cfg.freeze()

    set_seed(cfg.SOLVER.SEED)

    if cfg.MODEL.DIST_TRAIN:
        torch.cuda.set_device(args.local_rank)

    output_dir = cfg.OUTPUT_DIR
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    logger = setup_logger("EDITOR", output_dir, if_train=True)
    logger.info("Saving model in the path :{}".format(cfg.OUTPUT_DIR))
    logger.info(args)

    if args.config_file != "":
        logger.info("Loaded configuration file {}".format(args.config_file))
        with open(args.config_file, 'r') as cf:
            config_str = "\n" + cf.read()
            logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    if cfg.MODEL.DIST_TRAIN:
        torch.distributed.init_process_group(backend='nccl', init_method='env://')

    os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID

    model = make_model(cfg, num_class=50, camera_num=8)
    model.cuda()
    model.eval()
    print(sum(p.numel() / 1000000 for p in model.parameters() if p.requires_grad))



# Here is the code for testing the trainable parameters of the basic CNN model, different modality do not share the same parameters
# if __name__ == '__main__':
#     from modeling.backbones.basic_cnn_params import build_model as cnn_build_model
#     name = 'cal'
#     model = cnn_build_model(name, 50, use_gpu = True)
#     print("Now is ", name)
# Here the 3 means unshared parameters
#     print(3* sum(p.numel() / 1000000 for p in model.parameters() if p.requires_grad))