def load_id_name_map(cfg_path): config_map = get_config_map(cfg_path) file_path = config_map['id_name_txt'] mp = {} with open(file_path, 'r') as f: i = 0 for line in f: mp[i] = line.strip() i += 1 return mp
def load_classify_model(cfg_path, device='cuda:0'): config_map = get_config_map(cfg_path) model_ft, input_size = initialize_model(config_map['model_type'], config_map['class_number'], config_map['feature_extract'], use_pretrained=False) model_p = nn.DataParallel(model_ft.to(device), device_ids=config_map['gpu_ids']) model_p.load_state_dict(torch.load(config_map['load_from_path'])) model_p.eval() return model_p
import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler import torchvision from torch.autograd import Variable from torchvision import datasets, models, transforms import matplotlib.pyplot as plt from tqdm import tqdm from utils.train_utils import get_config_map from utils.model_utils import initialize_model # cfg_path = 'configs/classify2050c_densenet121_eval.json' cfg_path = 'configs/classify800c_se-resnext50_512_eval.json' config_map = get_config_map(cfg_path) input_resize = config_map['input_size'] # model_path = '/home/ubuntu/project/classify.pytorch/saved_models/densenet121_top100/epoch_3.pth' device = 'cuda:0' def load_model(model_path): print(model_path) model_ft = torch.load(model_path) num_ftrs = model_ft.classifier.in_features print(num_ftrs) input_size = 224 return model_ft, input_size