def load_model(self):
if self.use_dense:
model_structure = ft_net_dense(self.nclasses)
elif self.use_NAS:
model_structure = ft_net_NAS(self.nclasses)
else:
model_structure = ft_net(self.nclasses, stride=self.stride)
if self.PCB:
model_structure = PCB(self.nclasses)
#if self.fp16:
# model_structure = network_to_half(model_structure)
self.model = self.load_network(model_structure)
# Remove the final fc layer and classifier layer
if self.PCB:
#if self.fp16:
# model = PCB_test(model[1])
#else:
self.model = PCB_test(self.model)
else:
#if self.fp16:
#model[1].model.fc = nn.Sequential()
#model[1].classifier = nn.Sequential()
#else:
self.model.classifier.classifier = nn.Sequential()
# Change to test mode
self.model = self.model.eval()
if self.use_gpu:
self.model = self.model.cuda()
def get_model(class_names_size):
if opt.use_dense:
model = ft_net_dense(class_names_size, opt.droprate)
elif opt.use_NAS:
model = ft_net_NAS(class_names_size, opt.droprate)
else:
model = ft_net(class_names_size, opt.droprate, opt.stride)
if opt.PCB:
model = PCB(class_names_size)
return model
def loadModel(modelName, config, stateDictPath):
nclasses = config["nclasses"]
if modelName == "densenet121":
model_structure = ft_net_dense(nclasses)
elif modelName == "NAS":
model_structure = ft_net_NAS(nclasses)
elif modelName == "ResNet50":
model_structure = ft_net(nclasses, stride=config["stride"])
elif modelName == "PCB":
model_structure = PCB(nclasses)
else:
raise KeyError("Undefined modelName, %s" % modelName)
return loadNetwork(model_structure, stateDictPath)
gallery_cam, gallery_label = get_id(gallery_path)
query_cam, query_label = get_id(query_path)
train_all_cam, train_all_label = get_id(train_all_path)
if opt.multi:
mquery_path = image_datasets['multi-query'].imgs
mquery_cam, mquery_label = get_id(mquery_path)
######################################################################
# Load Collected data Trained model
print('-------test-----------')
if opt.use_dense:
model_structure = ft_net_dense(opt.nclasses)
elif opt.use_NAS:
model_structure = ft_net_NAS(opt.nclasses)
else:
model_structure = ft_net(opt.nclasses, stride=opt.stride)
if opt.PCB:
model_structure = PCB(opt.nclasses)
#if opt.fp16:
# model_structure = network_to_half(model_structure)
model = load_network(model_structure)
# Remove the final fc layer and classifier layer
if opt.PCB:
#if opt.fp16:
# model = PCB_test(model[1])
torch.save(network.cpu().state_dict(), save_path)
if torch.cuda.is_available():
network.cuda(gpu_ids[0])
######################################################################
# Finetuning the convnet
# ----------------------
#
# Load a pretrainied model and reset final fully connected layer.
#
if opt.use_dense:
model = ft_net_dense(len(class_names), opt.droprate)
elif opt.use_NAS:
model = ft_net_NAS(len(class_names), opt.droprate)
else:
model = ft_net(len(class_names), opt.droprate, opt.stride)
if opt.PCB:
model = PCB(len(class_names))
opt.nclasses = len(class_names)
print(model)
if not opt.PCB:
ignored_params = list(map(id, model.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer_ft = optim.SGD([{
def load_network(name, opt):
# Load config
dirname = os.path.join('./data/outputs', name)
last_model_name = os.path.basename(get_model_list(dirname, 'net'))
epoch = last_model_name.split('_')[1]
epoch = epoch.split('.')[0]
if not epoch == 'last':
epoch = int(epoch)
config_path = os.path.join(dirname, 'opts.yaml')
with open(config_path, 'r') as stream:
config = yaml.load(stream)
opt.name = config['name']
opt.inputsize = config['inputsize']
opt.data_dir = config['data_dir']
opt.train_all = config['train_all']
opt.train_veri = config['train_veri']
opt.train_comp = config['train_comp']
opt.train_comp_veri = config['train_comp_veri']
opt.droprate = config['droprate']
opt.color_jitter = config['color_jitter']
opt.batchsize = config['batchsize']
opt.inputsize = config['inputsize']
opt.stride = config['stride']
if 'pool' in config:
opt.pool = config['pool']
if 'use_DSE' in config:
opt.use_DSE = config['use_DSE']
if 'use_EF4' in config:
opt.use_EF4 = config['use_EF4']
opt.use_EF5 = config['use_EF5']
opt.use_EF6 = config['use_EF6']
if 'h' in config:
opt.h = config['h']
opt.w = config['w']
if 'gpu_ids' in config:
opt.gpu_ids = config['gpu_ids']
opt.erasing_p = config['erasing_p']
opt.lr = config['lr']
opt.nclasses = config['nclasses']
opt.erasing_p = config['erasing_p']
opt.use_dense = config['use_dense']
opt.use_NAS = config['use_NAS']
opt.use_SE = config['use_SE']
opt.use_IR = config['use_IR']
opt.PCB = config['PCB']
opt.CPB = config['CPB']
opt.fp16 = config['fp16']
opt.balance = config['balance']
opt.angle = config['angle']
opt.arc = config['arc']
if opt.use_dense:
model = ft_net_dense(opt.nclasses, opt.droprate, opt.stride, None,
opt.pool)
elif opt.use_NAS:
model = ft_net_NAS(opt.nclasses, opt.droprate, opt.stride)
elif opt.use_SE:
model = ft_net_SE(opt.nclasses, opt.droprate, opt.stride, opt.pool)
elif opt.use_DSE:
model = ft_net_DSE(opt.nclasses, opt.droprate, opt.stride, opt.pool)
elif opt.use_IR:
model = ft_net_IR(opt.nclasses, opt.droprate, opt.stride)
elif opt.use_EF4:
model = ft_net_EF4(opt.nclasses, opt.droprate)
elif opt.use_EF5:
model = ft_net_EF5(opt.nclasses, opt.droprate)
elif opt.use_EF6:
model = ft_net_EF6(opt.nclasses, opt.droprate)
else:
model = ft_net(opt.nclasses, opt.droprate, opt.stride, None, opt.pool)
if opt.PCB:
model = PCB(opt.nclasses)
if opt.CPB:
model = CPB(opt.nclasses)
if opt.angle:
model = ft_net_angle(opt.nclasses, opt.droprate, opt.stride)
elif opt.arc:
model = ft_net_arc(opt.nclasses, opt.droprate, opt.stride)
# load model
if isinstance(epoch, int):
save_filename = 'net_%03d.pth' % epoch
else:
save_filename = 'net_%s.pth' % epoch
save_path = os.path.join('./data/outputs', name, save_filename)
print('Load the model from %s' % save_path)
network = model
try:
network.load_state_dict(torch.load(save_path))
except:
network = torch.nn.DataParallel(network)
network.load_state_dict(torch.load(save_path))
network = network.module
return network, opt, epoch
