def main_train():
f = open(args.save_dir + 'train_log.txt', 'w')
copyfile('./train.py', args.save_dir + '/train.py')
copyfile('./model.py', args.save_dir + '/model.py')
model = PCB(len(datas["class"]))
if gpu:
model = model.cuda()
if is_parallel_train:
model = nn.DataParallel(model, device_ids=gpu_ids)
criterion = nn.CrossEntropyLoss()
model = pcb_train(model, criterion, f, "PCB", 60)
if args.RPP:
model = get_net(is_parallel_train, model).convert_to_rpp()
if use_gpu:
model = model.cuda()
if is_parallel_train:
model = nn.DataParallel(model, device_ids=gpu_ids)
model = rpp_train(model, criterion, f, "RPP", 5)
model = full_train(model, criterion, f, "full", 10)
f.close()
def get_reid(model_path):
checkpoint = torch.load(model_path)
model_structure = PCB(1453)
model_structure.load_state_dict(checkpoint)
model = PCB_test(model_structure)
model = model.eval()
return model
def main():
transform_test = T.Compose([
T.Resize((288,144),interpolation=3),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
print("I am in the main")
use_gpu = torch.cuda.is_available()
#model = torch.load(os.path.join('./model','best_model.pth.tar'))
#print(type(model))
#model.eval()
#model = models.init_model(name=args.arch, num_classes=751, loss={'xent'})
#if use_gpu:
# model = nn.DataParallel(model).cuda()
#checkpoint = torch.load(os.path.join('./model','net_last.pth'))
#save_path = os.path.join('./model','net_last.pth')
#model.load_state_dict(checkpoint['state_dict'])
#model.load_state_dict(torch.load(save_path))
#model = torch.load(os.path.join('./model','net_last.pth'))
#if use_gpu:
# model = nn.DataParallel(model).cuda()
#model.to(device)
#model_structure = models.init_model(name=args.arch, num_classes=751, loss={'xent'})
model_structure = PCB(751)
model=load_network(model_structure)
model = PCB_test(model)
model = model.eval()
if use_gpu:
model = model.cuda()
for subset in ['query', 'gallery']:
print(subset)
test_names, test_features = extractor(model, DataLoader(Dataset(subset,transform=transform_test)))
results = {'names': test_names, 'features': test_features.numpy()}
scipy.io.savemat(os.path.join('example', 'feature_val_%s.mat' % (subset)), results)
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 load_network(name, finetune=False):
# Load config
if finetune:
dirname = os.path.join('./model', 'ft_ResNet')
else:
dirname = os.path.join('./model', name)
last_model_name = os.path.basename(get_model_list(dirname, 'net'))
print(last_model_name)
epoch = last_model_name.split('_')[1]
epoch = epoch.split('.')[0]
if not epoch == 'last':
epoch = int(epoch)
if opt.PCB:
model = PCB(len(class_names), return_f=True)
else:
model = ft_net(len(class_names), return_f=True)
# load model
if isinstance(epoch, int):
save_filename = 'net_%03d.pth' % epoch
else:
save_filename = 'net_%s.pth' % epoch
if finetune:
save_path = os.path.join('./model', 'ft_ResNet', save_filename)
else:
save_path = os.path.join('./model', name, save_filename)
print('Load the model from %s' % save_path)
network = model
network.load_state_dict(torch.load(save_path))
return network, epoch
def __init__(self):
#base_model = VGG16(weights='imagenet')
#self.model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc1').output)
save_path = os.path.join('./model', 'net.pth')
model = PCB(575)
self.model = PCB_test(model, num_parts=6, cluster_plots=False)
self.model.load_state_dict(torch.load(save_path), strict=False)
self.avgpool = nn.AdaptiveAvgPool2d((3, 2))
self.model = self.model.eval()
def test_main():
print('-------test-----------')
gallery_path = images['gallery'].imgs
query_path = images['query'].imgs
gallery_cam, gallery_label = get_id(gallery_path)
query_cam, query_label = get_id(query_path)
model_structure = PCB(751)
if args.RPP:
model_structure = model_structure.convert_to_rpp()
model = load_network(model_structure)
model = PCB_test(model, feature_H)
model = model.eval()
if use_gpu:
model = model.cuda()
gallery_feature = extract_feature(model, datas['gallery'])
query_feature = extract_feature(model, datas['query'])
result = {
'gallery_feature': gallery_feature.numpy(),
'gallery_label': gallery_label,
'gallery_cam': gallery_cam,
'query_feature': query_feature.numpy(),
'query_label': query_label,
'query_cam': query_cam
}
if not os.path.isdir(args.result_dir):
os.mkdir(args.result_dir)
if args.RPP:
save_pre = '/RPP_'
else:
save_pre = '/PCB_'
if args.feature_H:
save_pre += 'H_'
else:
save_pre += 'G_'
scipy.io.savemat(args.result_dir + save_pre + 'result.mat', result)
def load_network(name, opt):
# Load config
dirname = os.path.join('./model', 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.data_dir = config['data_dir']
opt.train_all = config['train_all']
opt.droprate = config['droprate']
opt.color_jitter = config['color_jitter']
opt.batchsize = config['batchsize']
opt.h = config['h']
opt.w = config['w']
opt.stride = config['stride']
if 'pool' in config:
opt.pool = config['pool']
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.PCB = config['PCB']
opt.fp16 = config['fp16']
if opt.use_dense:
model = ft_net_dense(opt.nclasses, opt.droprate, opt.stride, None,
opt.pool)
else:
model = ft_net(opt.nclasses, opt.droprate, opt.stride, None, opt.pool)
if opt.PCB:
model = PCB(opt.nclasses)
# load model
if isinstance(epoch, int):
save_filename = 'net_%03d.pth' % epoch
else:
save_filename = 'net_%s.pth' % epoch
save_path = os.path.join('./model', name, save_filename)
print('Load the model from %s' % save_path)
network = model
network.load_state_dict(torch.load(save_path))
return network, opt, epoch
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)
def main():
print("==========\nArgs:{}\n==========".format(opt))
train_data_dir = opt.train_data_dir
test_data_dir = opt.test_data_dir
name = opt.name
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_ids
use_gpu = torch.cuda.is_available()
if opt.use_cpu: use_gpu = False
if not opt.evaluate:
sys.stdout = Logger(osp.join(opt.save_dir, opt.train_log))
else:
sys.stdout = Logger(osp.join(opt.save_dir, opt.test_log))
if use_gpu:
print("Currently using GPU {}".format(opt.gpu_ids))
cudnn.benchmark = True
else:
print("Currently using CPU (GPU is highly recommended)")
#str_ids = opt.gpu_ids.split(',')
#gpu_ids = []
#for str_id in str_ids:
# gid = int(str_id)
# if gid >=0:
# gpu_ids.append(gid)
#
## set gpu ids
#if len(gpu_ids)>0:
# torch.cuda.set_device(gpu_ids[0])
#print(gpu_ids[0])
# Load train Data
# ---------
print("==========Preparing trian dataset========")
transform_train_list = [
# transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC)
transforms.Resize((288, 144), interpolation=3),
transforms.RandomCrop((256, 128)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if opt.PCB:
transform_train_list = [
transforms.Resize((384, 192), interpolation=3),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if opt.erasing_p > 0:
transform_train_list = transform_train_list + [
RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0])
]
if opt.color_jitter:
transform_train_list = [
transforms.ColorJitter(
brightness=0.1, contrast=0.1, saturation=0.1, hue=0)
] + transform_train_list
train_data_transforms = transforms.Compose(transform_train_list)
train_all = ''
if opt.train_all:
if opt.use_clean_imgs:
train_all = '_all_clean'
else:
train_all = '_all'
print("Using all the train images")
train_image_datasets = {}
train_image_datasets['train'] = datasets.ImageFolder(
os.path.join(train_data_dir, 'train' + train_all),
train_data_transforms)
train_dataloaders = {
x: torch.utils.data.DataLoader(train_image_datasets[x],
batch_size=opt.train_batch,
shuffle=True,
num_workers=4)
for x in ['train']
}
dataset_sizes = {x: len(train_image_datasets[x]) for x in ['train']}
class_names = train_image_datasets['train'].classes
inputs, classes = next(iter(train_dataloaders['train']))
######################################################################
# Prepare test data
if not opt.train_only:
print("========Preparing test dataset========")
transform_test_list = transforms.Compose([
transforms.Resize((384, 192), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_image_datasets = {
x: datasets.ImageFolder(os.path.join(test_data_dir, x),
transform_test_list)
for x in ['gallery', 'query']
}
test_dataloaders = {
x: torch.utils.data.DataLoader(test_image_datasets[x],
batch_size=opt.test_batch,
shuffle=False,
num_workers=4)
for x in ['gallery', 'query']
}
print("Initializing model...")
if opt.use_dense:
model = ft_net_dense(len(class_names))
else:
model = ft_net(len(class_names))
if opt.PCB:
model = PCB(len(class_names))
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
start_epoch = opt.start_epoch
if opt.resume:
print("Loading checkpoint from '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
model.load_state_dict(checkpoint['state_dict'])
# model.load_state_dict(checkpoint)
start_epoch = checkpoint['epoch']
if use_gpu:
model = nn.DataParallel(model).cuda()
if opt.evaluate:
print("Evaluate only")
test(model, test_image_datasets, test_dataloaders, use_gpu)
return
criterion = nn.CrossEntropyLoss().cuda()
if opt.PCB:
ignored_params = list(map(id, model.module.resnet50.fc.parameters()))
ignored_params += (
list(map(id, model.module.classifier0.parameters())) +
list(map(id, model.module.classifier1.parameters())) +
list(map(id, model.module.classifier2.parameters())) +
list(map(id, model.module.classifier3.parameters())) +
list(map(id, model.module.classifier4.parameters())) +
list(map(id, model.module.classifier5.parameters()))
#+list(map(id, model.classifier7.parameters() ))
)
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer = optim.SGD(
[
{
'params': base_params,
'lr': 0.01
},
{
'params': model.module.resnet50.fc.parameters(),
'lr': 0.1
},
{
'params': model.module.classifier0.parameters(),
'lr': 0.1
},
{
'params': model.module.classifier1.parameters(),
'lr': 0.1
},
{
'params': model.module.classifier2.parameters(),
'lr': 0.1
},
{
'params': model.module.classifier3.parameters(),
'lr': 0.1
},
{
'params': model.module.classifier4.parameters(),
'lr': 0.1
},
{
'params': model.module.classifier5.parameters(),
'lr': 0.1
},
#{'params': model.classifier7.parameters(), 'lr': 0.01}
],
weight_decay=5e-4,
momentum=0.9,
nesterov=True)
else:
ignored_params = list(map(
id, model.module.model.fc.parameters())) + list(
map(id, model.module.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer = optim.SGD([{
'params': base_params,
'lr': 0.01
}, {
'params': model.module.model.fc.parameters(),
'lr': 0.1
}, {
'params': model.module.classifier.parameters(),
'lr': 0.1
}],
weight_decay=5e-4,
momentum=0.9,
nesterov=True)
# Decay LR by a factor of 0.1 every 40 epochs
if opt.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=opt.stepsize,
gamma=0.1)
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
######################################################################
# Training the model
# ------------------
#
# Now, let's write a general function to train a model. Here, we will
# illustrate:
#
# - Scheduling the learning rate
# - Saving the best model
#
# In the following, parameter ``scheduler`` is an LR scheduler object from
# ``torch.optim.lr_scheduler``.
for epoch in range(start_epoch, opt.max_epoch):
start_train_time = time.time()
if opt.train_only:
print("==> Training only")
train(epoch, model, criterion, optimizer, train_dataloaders,
use_gpu)
train_time += round(time.time() - start_train_time)
if epoch % opt.eval_step == 0:
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint(
{
'state_dict': state_dict,
'epoch': epoch,
}, 0,
osp.join(opt.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
if opt.stepsize > 0: scheduler.step()
else:
train(epoch, model, criterion, optimizer, train_dataloaders,
use_gpu)
train_time += round(time.time() - start_train_time)
if opt.stepsize > 0: scheduler.step()
if (epoch + 1) > opt.start_eval and opt.eval_step > 0 and (
epoch + 1) % opt.eval_step == 0 or (epoch +
1) == opt.max_epoch:
print("==> Test")
rank1 = test(model, test_image_datasets, test_dataloaders,
use_gpu)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint(
{
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, is_best,
osp.join(opt.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
if not opt.train_only:
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(
best_rank1, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
query_cam, query_label = get_id(query_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(751)
else:
model_structure = ft_net(751)
if opt.PCB:
model_structure = PCB(751)
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:
model = PCB_test(model)
model = network_to_half(model)
elif opt.fp16:
model[1].model.fc = nn.Sequential()
model[1].classifier = nn.Sequential()
else:
model.model.fc = nn.Sequential()
gallery_path = image_datasets['gallery'].imgs
query_path = image_datasets['query'].imgs
g_path = []
for gp in gallery_path:
(name, _) = gp
g_path.append(name.replace(' ', '').split('/')[-1])
q_path = []
for qp in query_path:
(name, _) = qp
q_path.append(name.replace(' ', '').split('/')[-1])
######################################################################
# Load Collected data Trained model
print('-------test-----------')
model_structure = PCB(class_num=4768) #for duke
if torch.cuda.device_count() > 1:
model_structure = nn.DataParallel(model_structure)
model = load_network(model_structure)
# Change to test mode
model = model.eval()
if use_gpu:
model = model.cuda()
# Extract feature
gallery_feature = extract_feature(model, dataloaders['gallery'])
query_feature = extract_feature(model, dataloaders['query'])
# Save to Matlab for check
######################################################################
# Finetuning the convnet
# ----------------------
#
# Load a pretrainied model and reset final fully connected layer.
#
if opt.use_dense:
model = ft_net_dense((id_class_number), opt.droprate)
else:
model = ft_net((id_class_number), attr_class_number, opt.droprate,
opt.stride)
if opt.PCB:
model = PCB((id_class_number))
opt.nclasses = (id_class_number)
print(model)
if not opt.PCB:
ignored_params = list(map(id, model.model.fc.parameters()))
for i in range(attr_class_number + 1):
list(map(id, model.__getattr__('class_' + str(i)).parameters()))
# ignored_params = list(map(id, model.model.fc.parameters())) + \
# list(map(id, model.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer_ft = optim.SGD(
mquery_path = image_datasets['multi-query'].imgs
mquery_cam, mquery_label = get_id(mquery_path)
######################################################################
# Load Collected data Trained model
print('-------test-----------')
class_number = {
'datasets/Market/datasets/pytorch/': 751,
'datasets/Duke/datasets/pytorch/': 702,
'datasets/CUHK03_detected/datasets/pytorch/': 767,
'datasets/CUHK03_labled/datasets/pytorch/': 767
}
class_num = class_number[data_dir]
part = 3
if not opt.mhn:
model_structure = PCB(class_num, part)
else:
parts = opt.parts
model_structure = MHN_smallPCB(class_num, parts, part)
model_structure.model.fc = nn.Sequential()
model_structure.model.avgpool = nn.Sequential()
#print(model_structure)
model = load_network(model_structure) ## load the learned params
if not opt.mhn:
model.model.avgpool = nn.Sequential()
model.model.fc = nn.Sequential()
# Remove the final fc layer and classifier layer
model.classifier0.classifier = nn.Sequential()
model.classifier1.classifier = nn.Sequential()
model.classifier2.classifier = nn.Sequential()
gallery_path = image_datasets['gallery'].imgs
query_path = image_datasets['query'].imgs
gallery_cam, gallery_label = get_id(gallery_path)
query_cam, query_label = get_id(query_path)
######################################################################
# Load Collected data Trained model
print('-------test-----------')
if opt.use_dense:
model_structure = ft_net_dense(28310)
else:
model_structure = ft_net(28310)
if opt.PCB:
model_structure = PCB(28310)
model_structure = nn.DataParallel(model_structure)
model = load_network(model_structure)
# Remove the final fc layer and classifier layer
#print(model)
if not opt.PCB:
model.module.model.fc = nn.Sequential()
# model.classifier = nn.Sequential()
else:
model = PCB_test(model)
print(model)
# Change to test mode
######################################################################
# Start training
# ---------------------------
# Record the training information #
f = open(args.save_dir + 'train_log.txt', 'w')
copyfile('./train.py', args.save_dir + '/train.py')
copyfile('./model.py', args.save_dir + '/model.py')
# step1: PCB training #
stage = 'PCB'
model = PCB(len(class_names))
if use_gpu:
model = model.cuda()
if is_parallel_train:
import pdb
#pdb.set_trace()
model = nn.DataParallel(model, device_ids=gpu_ids)
criterion = nn.CrossEntropyLoss()
model = pcb_train(model, criterion, f, stage, 100)
############################
# step2&3: RPP training #
if args.RPP:
print()
return model
# Save modeL
def save_network(network, epoch_label):
save_filename = 'net_%s.pth' % epoch_label
save_path = os.path.join(opt.outf, save_filename)
torch.save(network.cpu().state_dict(), save_path)
if torch.cuda.is_available:
network.cuda()
# setting of train phase
model = PCB(class_num=4768)
# load pretrained para without classifier
if use_gpu:
model = model.cuda()
# set the criterion
triplet_selector = SemihardNegativeTripletSelector(opt.margin)
criterion_tri = OnlineTripletLoss(opt.margin, triplet_selector)
criterion_part = nn.CrossEntropyLoss()
# criterion_part=CrossEntropyLabelSmooth(4768)
criterion_center = CenterLoss(4768)
criterion_focal = FocalLoss(gamma=2)
# updating rule for parameter
for i in range(len(dataset_list)):
dataset_path.append(image_datasets[dataset_list[i]].imgs)
dataset_cam = []
dataset_label = []
dataset_filename = []
for i in range(len(dataset_list)):
cam, label, filename = get_id(dataset_path[i])
dataset_cam.append(cam)
dataset_label.append(label)
dataset_filename.append(filename)
######################################################################
# Load Collected data Trained model
print('-------test-----------')
class_num = len(os.listdir(os.path.join(data_dir, 'train_all')))
model = PCB(class_num)
if 'st' in opt.which_epoch:
model = load_whole_network(model, name, opt.which_epoch + '_' + str(opt.net_loss_model))
else:
model = load_whole_network(model, name, opt.which_epoch)
model = model.eval()
if use_gpu:
model = model.cuda()
# Extract feature
dataset_feature = []
with torch.no_grad():
for i in range(len(dataset_list)):
dataset_feature.append(extract_feature(model, dataloaders[dataset_list[i]]))
result = {'gallery_f': dataset_feature[0].numpy(), 'gallery_label': dataset_label[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))
else:
model = ft_net(len(class_names))
if opt.PCB:
model = PCB(len(class_names))
print(model)
if use_gpu:
model = model.cuda()
criterion = nn.CrossEntropyLoss()
if not opt.PCB:
ignored_params = list(map(id, model.model.fc.parameters() )) + list(map(id, model.classifier.parameters() ))
base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
optimizer_ft = optim.SGD([
{'params': base_params, 'lr': 0.01},
{'params': model.model.fc.parameters(), 'lr': 0.1},
{'params': model.classifier.parameters(), 'lr': 0.1}
######################################################################
# Finetuning the convnet
# ----------------------
#
# Load a pretrainied model and reset final fully connected layer.
#
if opt.use_dense:
model = ft_net_dense(len(class_names))
else:
model = ft_net(len(class_names))
if opt.PCB:
model = PCB(len(class_names))
model_verif = verif_net()
# print(model)
# print(model_verif)
if use_gpu:
model = model.cuda()
model_verif = model_verif.cuda()
criterion = nn.CrossEntropyLoss()
if not opt.PCB:
ignored_params = list(map(id, model.model.fc.parameters())) + list(map(id, model.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
optimizer_ft = optim.SGD([
######################################################################
# 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([{
'params': base_params,
'lr': 0.1 * opt.lr
}, {
'params': model.classifier.parameters(),
'lr': opt.lr
else:
x_embedded = x
for i, t in enumerate(set(y)):
idx = y == t
plt.scatter(x_embedded[idx, 0], x_embedded[idx, 1], label=t)
plt.title(title, fontsize=16)
plt.legend()
if save is not None:
plt.savefig(save)
with open(CONFIG_FILE, 'r') as stream:
config = yaml.load(stream)
nclasses = config['nclasses']
model = PCB(nclasses)
model.load_state_dict(torch.load(WEIGHT_FILE))
model = model.eval().cuda()
features = []
labels = []
fol_list = os.listdir(SAMPLE_FOLDER)
for folder in fol_list:
if folder == 'ALL':
continue
fol = os.path.join(SAMPLE_FOLDER, folder)
if not os.path.isdir(fol):
continue
def test():
print('-------test-----------')
time_start = time.time()
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])
#else:
model = PCB_test(model)
else:
#if opt.fp16:
#model[1].model.fc = nn.Sequential()
#model[1].classifier = nn.Sequential()
#else:
model.classifier.classifier = nn.Sequential()
# Change to test mode
model = model.eval()
if use_gpu:
model = model.cuda()
# Extract feature
with torch.no_grad():
gallery_feature = extract_feature(model, dataloaders['gallery'])
query_feature = extract_feature(model, dataloaders['query'])
if opt.multi:
mquery_feature = extract_feature(model, dataloaders['multi-query'])
# Save to Matlab for check
result = {
'gallery_f': gallery_feature.numpy(),
'gallery_label': gallery_label,
'gallery_cam': gallery_cam,
'query_f': query_feature.numpy(),
'query_label': query_label,
'query_cam': query_cam
}
scipy.io.savemat('pytorch_result.mat', result)
print(opt.name)
os.system('python evaluate_gpu.py')
if opt.multi:
result = {
'mquery_f': mquery_feature.numpy(),
'mquery_label': mquery_label,
'mquery_cam': mquery_cam
}
scipy.io.savemat('multi_query.mat', result)
time_end = time.time()
print('totally cost', time_end - time_start)
query_path = image_datasets['query'].imgs
gallery_cam, gallery_label, gallery_frames = get_id(gallery_path)
query_cam, query_label, query_frames = get_id(query_path)
######################################################################
# Load Collected data Trained model
class_num = 751
print('-------test-----------')
if opt.use_dense:
model_structure = ft_net_dense(class_num)
else:
model_structure = ft_net(class_num)
if opt.PCB:
model_structure = PCB(class_num)
model = load_network(model_structure)
# Remove the final fc layer and classifier layer
if not opt.PCB:
model.model.fc = nn.Sequential()
model.classifier = nn.Sequential()
else:
model = PCB_test(model)
# Change to test mode
model = model.eval()
if use_gpu:
model = model.cuda()
def __init__(self, gpuID, model_path):
####################################################
# Options
# ------------------
self.gpu_ids = 1
self.which_epoch = '59'
self.batch_size = 1
self.use_dense = False
self.use_PCB = False
self.model_path = model_path
self.class_num = 751
self.score_threshold = 0.9
self.confidence_threshold = 0.6
####################################################
# Set gpu
# ------------------
use_gpu = torch.cuda.is_available()
if not use_gpu:
print('can not user gpu')
exit()
torch.cuda.set_device(self.gpu_ids)
####################################################
# Load model
# ------------------
print('load model...')
if self.use_dense:
model_structure = ft_net_dense(self.class_num)
else:
model_structure = ft_net(self.class_num)
if self.use_PCB:
model_structure = PCB(self.class_num)
model = utils.load_network(model_structure, self.model_path,
self.which_epoch)
# Remove the final fc layer and classifier layer
if not self.use_PCB:
model.model.fc = nn.Sequential()
model.classifier = nn.Sequential()
else:
model = PCB_test(model)
model = model.eval()
#print(model)
if use_gpu:
model = model.cuda()
self.model = model
####################################################
# Set Transform
# ------------------
print('set transform...')
self.data_transforms = transforms.Compose([
transforms.Resize((256, 128), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
if self.use_PCB:
self.data_transforms = transforms.Compose([
transforms.Resize((384, 192), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
######################################################################
# Finetuning the convnet
# ----------------------
#
# Load a pretrainied model and reset final fully connected layer.
#
if opt.use_dense:
model = ft_net_dense(len(class_names))
else:
#model = ft_net(len(class_names))
model = ResidualAttentionModel(len(class_names))
if opt.PCB:
model = PCB(len(class_names))
print(model)
if use_gpu:
model = model.cuda()
criterion = nn.CrossEntropyLoss()
lr = 0.001
optimizer_ft = torch.optim.Adam(model.parameters(), lr=lr)
if not opt.PCB:
print("Here")
#ignored_params = list(map(id, model.fc.parameters() )) + list(map(id, model.classifier.parameters() ))
#base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
#optimizer_ft = optim.SGD([
# {'params': base_params, 'lr': 0.01},
labels.append(int(label))
camera_id.append(int(camera[0]))
return camera_id, labels
gallery_path = image_datasets['gallery'].imgs
query_path = image_datasets['query'].imgs
gallery_cam, gallery_label = get_id(gallery_path)
query_cam, query_label = get_id(query_path)
######################################################################
# Load Collected data Trained model
print('-------test-----------')
model_structure = PCB(751)
if args.RPP:
model_structure = model_structure.convert_to_rpp()
model = load_network(model_structure)
model = PCB_test(model, feature_H)
# Change to test mode
model = model.eval()
if use_gpu:
model = model.cuda()
# Extract feature
gallery_feature = extract_feature(model, dataloaders['gallery'])
query_feature = extract_feature(model, dataloaders['query'])
def train(opt):
version = torch.__version__
fp16 = opt.fp16
data_dir = opt.data_dir
name = opt.name
str_ids = opt.gpu_ids.split(',')
gpu_ids = []
for str_id in str_ids:
gid = int(str_id)
if gid >= 0:
gpu_ids.append(gid)
# set gpu ids
if len(gpu_ids) > 0:
torch.cuda.set_device(gpu_ids[0])
cudnn.benchmark = True
######################################################################
# Load Data
# ---------
#
transform_train_list = [
# transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC)
transforms.Resize((256, 128), interpolation=3),
transforms.Pad(10),
transforms.RandomCrop((256, 128)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
transform_val_list = [
transforms.Resize(size=(256, 128), interpolation=3), # Image.BICUBIC
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if opt.PCB:
transform_train_list = [
transforms.Resize((384, 192), interpolation=3),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
transform_val_list = [
transforms.Resize(size=(384, 192),
interpolation=3), # Image.BICUBIC
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if opt.erasing_p > 0:
transform_train_list = transform_train_list + [
RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0])
]
if opt.color_jitter:
transform_train_list = [
transforms.ColorJitter(
brightness=0.1, contrast=0.1, saturation=0.1, hue=0)
] + transform_train_list
# print(transform_train_list)
data_transforms = {
'train': transforms.Compose(transform_train_list),
'val': transforms.Compose(transform_val_list),
}
train_all = ''
if opt.train_all:
train_all = '_all'
image_datasets = {}
image_datasets['train'] = datasets.ImageFolder(
os.path.join(data_dir, 'train' + train_all), data_transforms['train'])
image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'),
data_transforms['val'])
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=opt.batchsize,
shuffle=True,
num_workers=8,
pin_memory=True)
# 8 workers may work faster
for x in ['train', 'val']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
use_gpu = torch.cuda.is_available()
#since = time.time()
#inputs, classes = next(iter(dataloaders['train']))
#print('time used for loading data: %ds' %(time.time() - since))
######################################################################
# Training the model
# ------------------
#
# Now, let's write a general function to train a model. Here, we will
# illustrate:
#
# - Scheduling the learning rate
# - Saving the best model
#
# In the following, parameter ``scheduler`` is an LR scheduler object from
# ``torch.optim.lr_scheduler``.
y_loss = {} # loss history
y_loss['train'] = []
y_loss['val'] = []
y_err = {}
y_err['train'] = []
y_err['val'] = []
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
since = time.time()
results = []
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
scheduler.step()
model.train(True) # Set model to training mode
else:
model.train(False) # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0.0
# Iterate over data.
pbar = tqdm(dataloaders[phase])
for inputs, labels in pbar:
# get the inputs
now_batch_size, c, h, w = inputs.shape
if now_batch_size < opt.batchsize: # skip the last batch
continue
# print(inputs.shape)
# wrap them in Variable
if use_gpu:
inputs = Variable(inputs.cuda().detach())
labels = Variable(labels.cuda().detach())
else:
inputs, labels = Variable(inputs), Variable(labels)
# if we use low precision, input also need to be fp16
# if fp16:
# inputs = inputs.half()
# zero the parameter gradients
optimizer.zero_grad()
# forward
if phase == 'val':
with torch.no_grad():
outputs = model(inputs)
else:
outputs = model(inputs)
if not opt.PCB:
_, preds = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)
else:
part = {}
sm = nn.Softmax(dim=1)
num_part = 6
for i in range(num_part):
part[i] = outputs[i]
score = sm(part[0]) + sm(part[1]) + sm(part[2]) + sm(
part[3]) + sm(part[4]) + sm(part[5])
_, preds = torch.max(score.data, 1)
loss = criterion(part[0], labels)
for i in range(num_part - 1):
loss += criterion(part[i + 1], labels)
# backward + optimize only if in training phase
if phase == 'train':
if fp16: # we use optimier to backward loss
with amp.scale_loss(loss,
optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
# statistics
if int(version[0]) > 0 or int(
version[2]
) > 3: # for the new version like 0.4.0, 0.5.0 and 1.0.0
running_loss += loss.item() * now_batch_size
else: # for the old version like 0.3.0 and 0.3.1
running_loss += loss.data[0] * now_batch_size
running_corrects += float(torch.sum(preds == labels.data))
pbar.set_description(
desc='loss: {:.4f}'.format(loss.item()))
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects / dataset_sizes[phase]
print('\r\n{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
logging.info('epoch: {}, {} Loss: {:.4f} Acc: {:.4f}'.format(
epoch, phase, epoch_loss, epoch_acc))
y_loss[phase].append(epoch_loss)
y_err[phase].append(1.0 - epoch_acc)
# deep copy the model
if phase == 'val':
results.append({
'epoch': epoch,
'trainLoss': y_loss['train'][-1],
'trainError': y_err['train'][-1],
'valLoss': y_loss['val'][-1],
'valError': y_err['val'][-1]
})
last_model_wts = model.state_dict()
if epoch % 10 == 9:
save_network(model, epoch)
draw_curve(epoch)
write_to_csv(results)
time_elapsed = time.time() - since
print('\r\nTraining complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print()
time_elapsed = time.time() - since
print('\r\nTraining complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
model.load_state_dict(last_model_wts)
save_network(model, 'last')
return model
######################################################################
# Draw Curve
# ---------------------------
x_epoch = []
fig = plt.figure()
ax0 = fig.add_subplot(121, title="loss")
ax1 = fig.add_subplot(122, title="top1err")
def draw_curve(current_epoch):
x_epoch.append(current_epoch)
ax0.plot(x_epoch, y_loss['train'], 'bo-', label='train')
ax0.plot(x_epoch, y_loss['val'], 'ro-', label='val')
ax1.plot(x_epoch, y_err['train'], 'bo-', label='train')
ax1.plot(x_epoch, y_err['val'], 'ro-', label='val')
if current_epoch == 0:
ax0.legend()
ax1.legend()
fig.savefig(os.path.join('./model', name, 'train.jpg'))
def write_to_csv(results):
path = os.path.join('./model', name, 'result.csv')
with open(path, 'w', newline='') as csvfile:
writer = csv.DictWriter(csvfile,
fieldnames=list(results[0].keys()))
writer.writeheader()
writer.writerows(results)
######################################################################
# Save model
# ---------------------------
def save_network(network, epoch_label):
save_filename = 'net_%s.pth' % epoch_label
rpth = os.path.join('./model', name, 'Model Files')
if not os.path.exists(rpth):
os.makedirs(rpth)
save_path = os.path.join(rpth, save_filename)
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)
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)
print('model loaded')
if not opt.PCB:
ignored_params = list(map(id, model.model.fc.parameters())) + list(
map(id, model.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer_ft = optim.SGD([{
'params': base_params,
'lr': 0.1 * opt.lr
}, {
'params': model.model.fc.parameters(),
'lr': opt.lr
}, {
'params': model.classifier.parameters(),
'lr': opt.lr
}],
weight_decay=5e-4,
momentum=0.9,
nesterov=True)
else:
ignored_params = list(map(id, model.model.fc.parameters()))
ignored_params += (
list(map(id, model.classifier0.parameters())) +
list(map(id, model.classifier1.parameters())) +
list(map(id, model.classifier2.parameters())) +
list(map(id, model.classifier3.parameters())) +
list(map(id, model.classifier4.parameters())) +
list(map(id, model.classifier5.parameters()))
# +list(map(id, model.classifier6.parameters() ))
# +list(map(id, model.classifier7.parameters() ))
)
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer_ft = optim.SGD(
[
{
'params': base_params,
'lr': 0.1 * opt.lr
},
{
'params': model.model.fc.parameters(),
'lr': opt.lr
},
{
'params': model.classifier0.parameters(),
'lr': opt.lr
},
{
'params': model.classifier1.parameters(),
'lr': opt.lr
},
{
'params': model.classifier2.parameters(),
'lr': opt.lr
},
{
'params': model.classifier3.parameters(),
'lr': opt.lr
},
{
'params': model.classifier4.parameters(),
'lr': opt.lr
},
{
'params': model.classifier5.parameters(),
'lr': opt.lr
},
# {'params': model.classifier6.parameters(), 'lr': 0.01},
# {'params': model.classifier7.parameters(), 'lr': 0.01}
],
weight_decay=5e-4,
momentum=0.9,
nesterov=True)
# Decay LR by a factor of 0.1 every 40 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
step_size=40,
gamma=0.1)
######################################################################
# Train and evaluate
# ^^^^^^^^^^^^^^^^^^
#
# It should take around 1-2 hours on GPU.
#
dir_name = os.path.join('./model', name)
if not os.path.isdir(dir_name):
os.mkdir(dir_name)
# record every run
copyfile('./train.py', dir_name + '/train.py')
copyfile('./model.py', dir_name + '/model.py')
# save opts
with open('%s/opts.yaml' % dir_name, 'w') as fp:
yaml.dump(vars(opt), fp, default_flow_style=False)
# model to gpu
model = model.cuda()
if fp16:
# model = network_to_half(model)
# optimizer_ft = FP16_Optimizer(optimizer_ft, static_loss_scale = 128.0)
model, optimizer_ft = amp.initialize(model,
optimizer_ft,
opt_level="O1")
criterion = losses.DualLoss()
model = train_model(model,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=60)
#
# if __name__ == "__main__":
# train(opt)
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])
#else:
model = PCB_test(model)
else:
#if opt.fp16:
#model[1].model.fc = nn.Sequential()
save_whole_network(model, name, 'last' + '_' + str(opt.net_loss_model))
return model
######################################################################
# Train and evaluate
# ^^^^^^^^^^^^^^^^^^
#
# It should take around 1-2 hours on GPU.
#
dir_name = os.path.join('./model', name)
if not os.path.exists('model'):
os.mkdir('model')
print('class_num = %d' % (class_num))
model = PCB(class_num, train=True)
if use_gpu:
model.cuda()
# print('model structure')
# print(model)
criterion = nn.CrossEntropyLoss()
criterion_soft = SoftLabelLoss()
classifier_id = (list(map(id, model.classifier0.parameters())) +
list(map(id, model.classifier1.parameters())) +
list(map(id, model.classifier2.parameters())) +
list(map(id, model.classifier3.parameters())) +
list(map(id, model.classifier4.parameters())) +
list(map(id, model.classifier5.parameters())))
classifier_params = filter(lambda p: id(p) in classifier_id,
