nargs='+',
type=int,
default=[10, 20, 50],
help='ks for Metric@k')
parser.add_argument('--best_metric',
type=str,
default='NDCG@10',
help='Metric for determining the best model')
################
# Model
################
parser.add_argument('--model_code',
type=str,
default='bert',
choices=MODELS.keys())
parser.add_argument('--model_init_seed', type=int, default=None)
# BERT #
parser.add_argument('--bert_max_len',
type=int,
default=None,
help='Length of sequence for bert')
parser.add_argument('--bert_num_items',
type=int,
default=None,
help='Number of total items')
parser.add_argument('--bert_hidden_units',
type=int,
default=None,
help='Size of hidden vectors (d_model)')
parser.add_argument('--bert_num_blocks',
import pytorch_retinanet.model_pnasnet
import pytorch_retinanet.model_resnet
import pytorch_retinanet.model_se_resnext
import pytorch_retinanet.model_xception
import torch
from config import DATA_DIR, IMG_SIZE, RESULTS_DIR, TEST_DIR, WEIGHTS_DIR
from datasets.test_dataset import TestDataset
from models import MODELS
from torch import nn, optim
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from torchvision import datasets, models, transforms
from utils.logger import Logger
from utils.my_utils import set_seed
model_configs = MODELS.keys()
def load_model(checkpoint: str) -> nn.Module:
"""
Helper to load model weihts
"""
print(f"Loading model from: {checkpoint}")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# load model
model = torch.load(checkpoint)
model = model.to(device)
model.eval()
# model = torch.nn.DataParallel(model).cuda()
return model