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