device = torch.device("cuda" if (
torch.cuda.is_available() and CONFIG.ngpu > 0) else "cpu")
else:
device = torch.device("cpu")
get_logger(CONFIG.log_dir)
writer = get_writer(args.title, CONFIG.write_dir)
logging.info(
"=================================== Experiment title : {} Start ==========================="
.format(args.title))
set_random_seed(CONFIG.seed)
train_transform, val_transform, test_transform = get_transforms(CONFIG)
train_dataset, val_dataset, test_dataset = get_dataset(
train_transform, val_transform, test_transform, CONFIG)
train_loader, val_loader, test_loader = get_dataloader(
train_dataset, val_dataset, test_dataset, CONFIG)
generator = get_generator(CONFIG, 21 * 8)
generator.to(device)
# ============ OFA ================
accuracy_predictor = AccuracyPredictor(pretrained=True, device=device)
print(accuracy_predictor.model)
flops_table = FLOPsTable(device=device)
# =================================
g_optimizer = get_optimizer(generator, CONFIG.g_optim_state)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--cfg", type=str, help="path to the config file", required=True)
parser.add_argument("--model", type=str, help="model", required=True)
parser.add_argument("--pretrain", type=str, help="path to the pretrain weight", required=True)
args = parser.parse_args()
CONFIG = get_config(args.cfg)
device = torch.device("cuda") if torch.cuda.is_available() and CONFIG.cuda else torch.device("cpu")
logger = get_logger(CONFIG.log_dir)
_, _, test_transform = get_transforms(CONFIG)
_, _, test_dataset = get_dataset(None, None, test_transform, CONFIG)
_, _, test_loader = get_dataloader(None, None, test_dataset, CONFIG)
if args.model == "PONASA":
model = PONASA()
elif args.model == "PONASB":
model = PONASB()
elif args.model == "PONASC":
model = PONASC()
model.load_state_dict(torch.load(args.pretrain))
latency = calculate_latency(model, 3, CONFIG.input_size)
counter = MAC_Counter(model, [1, 3, CONFIG.input_size, CONFIG.input_size])
macs = counter.print_summary(False)
param_nums = calculate_param_nums(model)
