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