def transform(self):
transform_train = [
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 = [
transforms.Resize(size=(256, 128), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if self.erasing_p > 0:
transform_train = transform_train + [
RandomErasing(probability=self.erasing_p, mean=[0.0, 0.0, 0.0])
]
if self.color_jitter:
transform_train = [
transforms.ColorJitter(
brightness=0.1, contrast=0.1, saturation=0.1, hue=0)
] + transform_train
self.data_transforms = {
'train': transforms.Compose(transform_train),
'val': transforms.Compose(transform_val),
}
def __init__(self, args):
train_list = [
transforms.Resize((args.height, args.width), interpolation=3),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
if args.random_erasing:
train_list.append(
RandomErasing(probability=args.probability,
mean=[0.0, 0.0, 0.0]))
train_transform = transforms.Compose(train_list)
test_transform = transforms.Compose([
transforms.Resize((args.height, args.width), interpolation=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
if not args.test_only:
module_train = import_module('data.' + args.data_train.lower())
self.trainset = getattr(module_train,
args.data_train)(args, train_transform,
'train')
self.train_loader = dataloader.DataLoader(
self.trainset,
sampler=RandomSampler(self.trainset,
args.batchid,
batch_image=args.batchimage),
#shuffle=True,
batch_size=args.batchid * args.batchimage,
num_workers=args.nThread)
else:
self.train_loader = None
if args.data_test in ['Market1501']:
module = import_module('data.' + args.data_train.lower())
self.testset = getattr(module,
args.data_test)(args, test_transform,
'test')
self.queryset = getattr(module,
args.data_test)(args, test_transform,
'query')
else:
raise Exception()
self.test_loader = dataloader.DataLoader(self.testset,
batch_size=args.batchtest,
num_workers=args.nThread)
self.query_loader = dataloader.DataLoader(self.queryset,
batch_size=args.batchtest,
num_workers=args.nThread)
def __init__(self, data_cfg, multi=1, nl=False):
"""
Dataset for training.
:param data_cfg: CfgNode for CityFlow NL.
"""
self.nl = nl
self.multi = multi
self.motion = data_cfg.motion
self.nseg = data_cfg.nseg
self.all3 = data_cfg.all3
self.pad = data_cfg.pad
self.data_cfg = data_cfg
self.aug = AutoAugment(auto_augment_policy(name='v0r', hparams=None))
with open(self.data_cfg.JSON_PATH) as f:
tracks = json.load(f)
f.close()
self.list_of_uuids = list(tracks.keys())
self.list_of_tracks = list(tracks.values())
self.list_of_crops = list()
train_num = len(self.list_of_uuids)
self.transform = transforms.Compose([
transforms.Pad(10),
transforms.RandomCrop(
(data_cfg.CROP_SIZE, self.data_cfg.CROP_SIZE)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
RandomErasing(probability=0.5)
])
if data_cfg.semi:
#cv
with open(self.data_cfg.EVAL_TRACKS_JSON_PATH) as f:
unlabel_tracks = json.load(f)
f.close()
self.list_of_uuids.extend(unlabel_tracks.keys())
self.list_of_tracks.extend(unlabel_tracks.values())
#nl
with open("data/test-queries.json", "r") as f:
unlabel_nl = json.load(f)
unlabel_nl_key = list(unlabel_nl.keys())
print('#track id (class): %d ' % len(self.list_of_tracks))
count = 0
# add id and nl, -1 for unlabeled data
for track_idx, track in enumerate(self.list_of_tracks):
track["track_id"] = track_idx
track["nl_id"] = track_idx
# from 0 to train_num-1 is the id of the original training set.
if track_idx >= train_num:
track["nl_id"] = -1
track["nl"] = unlabel_nl[unlabel_nl_key[count]]
count = count + 1
self._logger = get_logger()
def __init__(self, data_pth, is_train=True, *args, **kwargs):
super(Market1501, self).__init__(*args, **kwargs)
## parse image names to generate image ids
imgs = os.listdir(data_pth)
imgs = [im for im in imgs if osp.splitext(im)[-1] == '.jpg']
self.is_train = is_train
self.im_pths = [osp.join(data_pth, im) for im in imgs]
self.im_infos = {}
self.person_infos = {}
for i, im in enumerate(imgs):
tokens = im.split('_')
im_pth = self.im_pths[i]
pid = int(tokens[0])
cam = int(tokens[1][1])
self.im_infos.update({im_pth: (pid, cam)})
if pid in self.person_infos.keys():
self.person_infos[pid].append(i)
else:
self.person_infos[pid] = [
i,
]
self.pid_label_map = {}
for i, (pid, ids) in enumerate(self.person_infos.items()):
self.person_infos[pid] = np.array(ids, dtype=np.int32)
self.pid_label_map[pid] = i
## preprocessing
self.trans_train = transforms.Compose([
transforms.Resize((288, 144)),
transforms.RandomCrop((256, 128)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
RandomErasing(0.5, mean=[0.0, 0.0, 0.0])
])
## H-Flip
self.trans_no_train_flip = transforms.Compose([
transforms.Resize((288, 144)),
transforms.RandomHorizontalFlip(1),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
self.trans_no_train_noflip = transforms.Compose([
transforms.Resize((288, 144)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
def __init__(self,
crop_prob=0,
crop_ratio=1.0,
resize_h_w=None,
scale=True,
im_mean=None,
im_std=None,
mirror_type=None,
batch_dims='NCHW',
is_random_erasing=False,
prng=np.random):
"""
Args:
crop_prob: the probability of each image to go through cropping
crop_ratio: a float. If == 1.0, no cropping.
resize_h_w: (height, width) after resizing. If `None`, no resizing.
scale: whether to scale the pixel value by 1/255
im_mean: (Optionally) subtracting image mean; `None` or a tuple or list or
numpy array with shape [3]
im_std: (Optionally) divided by image std; `None` or a tuple or list or
numpy array with shape [3]. Dividing is applied only when subtracting
mean is applied.
mirror_type: How image should be mirrored; one of
[None, 'random', 'always']
batch_dims: either 'NCHW' or 'NHWC'. 'N': batch size, 'C': num channels,
'H': im height, 'W': im width. PyTorch uses 'NCHW', while TensorFlow
uses 'NHWC'.
prng: can be set to a numpy.random.RandomState object, in order to have
random seed independent from the global one
"""
self.crop_prob = crop_prob
self.crop_ratio = crop_ratio
self.resize_h_w = resize_h_w
self.scale = scale
self.im_mean = im_mean
self.im_std = im_std
self.check_mirror_type(mirror_type)
self.mirror_type = mirror_type
self.check_batch_dims(batch_dims)
self.batch_dims = batch_dims
self.prng = prng
self.is_random_erasing = is_random_erasing
self.RandomErasing = RandomErasing(probability=0.5,
sh=0.2,
mean=[0.0, 0.0, 0.0])
def __init__(self, pids_list, array, is_train=True, *args, **kwargs):
super(Market1501, self).__init__(*args, **kwargs)
self.is_train = is_train
self.pids = pids_list
self.array = array
self.img_names = [
el.split("/")[1] for el in self.pids
if os.path.splitext(el)[1] == '.jpg'
]
self.lb_ids = [int(el.split('_')[0]) for el in self.img_names]
self.lb_cams = [int(el.split('_')[1][1]) for el in self.img_names]
if is_train:
self.trans = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((288, 144)),
transforms.RandomCrop((256, 128)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.486, 0.459, 0.408),
(0.229, 0.224, 0.225)),
RandomErasing(0.5, mean=[0.0, 0.0, 0.0])
])
else:
self.trans_tuple = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.486, 0.459, 0.408),
(0.229, 0.224, 0.225))
])
self.Lambda = transforms.Lambda(
lambda crops: [self.trans_tuple(crop) for crop in crops])
self.trans = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((288, 144)),
transforms.TenCrop((256, 128)), self.Lambda
])
# useful for sampler
self.lb_img_dict = dict()
self.lb_ids_uniq = set(self.lb_ids)
lb_array = np.array(self.lb_ids)
for lb in self.lb_ids_uniq:
idx = np.where(lb_array == lb)[0]
self.lb_img_dict.update({lb: idx})
######################################################################
# 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(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.erasing_p>0:
transform_train_list = transform_train_list + [RandomErasing(opt.erasing_p)]
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)
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])
]
data_transforms = {
'train': transforms.Compose( transform_train_list ),
def main():
transforms_args = [
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
train_dataset = CoCoDataset(
args.coco_path,
"training",
target_size=args.target_size,
transform=transforms.Compose(
transforms_args +
[RandomErasing(probability=args.p, sh=args.sh, r1=args.r1)]))
test_dataset = CoCoDataset(args.coco_path,
"validation_wo_occlusion",
target_size=args.target_size,
transform=transforms.Compose(transforms_args))
train_batch_sampler = TrainBalancedBatchSampler(torch.from_numpy(
np.array(train_dataset.all_targets())),
K=args.K,
P=args.P,
n_batches=args.n_batches)
test_batch_sampler = TestBalancedBatchSampler(torch.from_numpy(
np.array(test_dataset.all_targets())),
K=args.K,
P=args.P,
n_batches=args.n_batches)
train_loader = DataLoader(train_dataset,
batch_sampler=train_batch_sampler,
**kwargs)
test_loader = DataLoader(test_dataset,
batch_sampler=test_batch_sampler,
**kwargs)
# init model
model, optim_state_dict, init_epoch = load_model(
args.backbone,
args.snapshot,
imagenet_weights=args.imagenet_weights,
freeze=args.freeze)
print("Resume training from epoch", init_epoch)
if cuda:
model.cuda()
# init optimizer
if args.optim == "Adam":
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=1e-4)
elif args.optim == "SGD":
optimizer = optim.SGD(model.parameters(),
momentum=0.9,
lr=args.lr,
weight_decay=1e-4)
else:
raise ValueError("Optimizer is not supported")
if optim_state_dict is not None:
optimizer.load_state_dict(optim_state_dict)
# define loss function
if args.triplet_selector == "hard":
selector = HardestNegativeTripletSelector(args.soft_margin)
elif args.triplet_selector == "semi":
selector = SemihardNegativeTripletSelector(args.soft_margin)
elif args.triplet_selector == "random":
selector = RandomNegativeTripletSelector(args.soft_margin)
else:
selector = AllTripletSelector()
train_loss_fn = TripletLoss(selector, soft_margin=args.soft_margin)
test_loss_fn = TripletLoss(AllTripletSelector(),
soft_margin=args.soft_margin)
# define learning rate scheduler
lr_scheduler = LrScheduler(args.epoch_decay_start, args.n_epoch, args.lr)
log_file = os.path.join(
args.logger_dir, '%s_%s.csv' % (args.backbone, args.triplet_selector))
for epoch in range(init_epoch + 1, args.n_epoch):
lr_scheduler.adjust_learning_rate(optimizer, epoch, args.optim)
for param_group in optimizer.param_groups:
print("LR: ", param_group['lr'])
train_loss = train_epoch(model, train_loader, train_loss_fn, optimizer,
cuda)
if epoch % args.eval_freq == 0:
test_loss = test_epoch(model, test_loader, test_loss_fn, cuda)
print('Epoch [%d/%d], Train loss: %.4f, Test loss: %.4f' %
(epoch, args.n_epoch, train_loss, test_loss))
log = [epoch, train_loss, test_loss]
if os.path.isfile(log_file):
with open(log_file, mode='a', newline='') as csv_f:
writer = csv.writer(csv_f)
writer.writerow(log)
else:
with open(log_file, mode='w', newline='') as csv_f:
writer = csv.writer(csv_f)
# write header
writer.writerow(["epoch", "train_loss", "test_loss"])
writer.writerow(log)
if epoch % args.save_freq == 0:
torch.save(
{
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'epoch': epoch
},
os.path.join(
args.snapshot_path, '%s_%s_%d.pth' %
(args.backbone, args.triplet_selector, epoch)))
transforms.RandomHorizontalFlip(0.5),
transforms.Pad(10),
transforms.RandomCrop([288, 144]),
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 erasing_p > 0:
transform_train_list = transform_train_list + [
RandomErasing(probability=erasing_p)
]
print(transform_train_list)
data_transforms = {
'train': transforms.Compose(transform_train_list),
'val': transforms.Compose(transform_val_list),
}
train_all = ''
if train_all_1:
train_all = '_all'
image_datasets = {}
image_datasets['train'] = datasets.ImageFolder(
os.path.join(data_dir, 'train' + train_all), data_transforms['train'])
def main():
global best_rank1
global pn_gan
global poses
if args.use_gan:
pn_gan, poses = initialize_PN_GAN()
torch.manual_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(
root=args.root,
name=args.dataset,
split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled,
cuhk03_classic_split=args.cuhk03_classic_split,
)
transform_train_list = ([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
if args.erasing_p > 0:
transform_train_list = transform_train_list + [
RandomErasing(probability=args.erasing_p, mean=[0.0, 0.0, 0.0])
]
transform_train = T.Compose(transform_train_list)
pin_memory = True if use_gpu else False
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent'},
use_gpu=use_gpu)
print("Model size: {:.3f} M".format(count_num_param(model)))
if args.label_smooth:
criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids,
use_gpu=use_gpu)
else:
criterion = nn.CrossEntropyLoss()
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
if args.fixbase_epoch > 0:
if hasattr(model, 'classifier') and isinstance(model.classifier,
nn.Module):
optimizer_tmp = init_optim(args.optim,
model.classifier.parameters(),
args.fixbase_lr, args.weight_decay)
else:
print(
"Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0"
)
args.fixbase_epoch = 0
if args.load_weights and check_isfile(args.load_weights):
# load pretrained weights but ignore layers that don't match in size
model = torch.load('')
model.eval()
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if args.resume and check_isfile(args.resume):
checkpoint_saved = torch.load('./checkpoint_self_after_reset8.pth.tar')
model.load_state_dict(checkpoint_saved)
model.train()
print("Loaded checkpoint from '{}'".format(args.resume))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
distmat = test(model,
queryloader,
galleryloader,
use_gpu,
return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(
distmat,
dataset,
save_dir=osp.join(args.save_dir, 'ranked_results'),
topk=20,
)
return
start_time = time.time()
train_time = 0
best_epoch = args.start_epoch
print("==> Start training")
if args.fixbase_epoch > 0:
print(
"Train classifier for {} epochs while keeping base network frozen".
format(args.fixbase_epoch))
for epoch in range(args.fixbase_epoch):
start_train_time = time.time()
train(epoch,
model,
criterion,
optimizer_tmp,
trainloader,
use_gpu,
freeze_bn=True)
train_time += round(time.time() - start_train_time)
del optimizer_tmp
print("Now open all layers for training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(state_dict,
'checkpoint_self_after_reset' + str(epoch + 1) + '.pth.tar')
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, 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(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
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))
def __init__(self,
data_path,
train_file,
img_size,
is_train=True,
*args,
**kwargs):
super(VehicleID, self).__init__(*args, **kwargs)
self.is_train = is_train
self.data_path = data_path
self.check = os.listdir(data_path)
self.check = [
el for el in self.check if os.path.splitext(el)[1] == '.jpg'
]
self.dict = {}
f = open(train_file, 'r')
self.train_file = f.readlines()
self.imgs = []
max_class = 0
self.classes = set()
now = 0
for line in self.train_file:
line = line.strip()
if (len(line.split(' ')) == 2):
car_name, car_class = line.split(' ')[0], int(
line.split(' ')[1])
elif (len(line.split(' ')) == 1 and not is_train):
car_name, car_class = line.split(' ')[0], -1
else:
logger.info('dataset wrong')
#if car_name not in self.check:
# logger.warning("%s do not exists"%(car_name))
# continue
if (len(car_name.split('.')) == 1):
car_name = car_name + '.jpg'
if car_name not in self.dict:
self.dict[car_name] = car_class
max_class = max(max_class, int(car_class))
self.classes.add(car_class)
self.imgs.append(car_name)
logger.info('dataset OK')
if self.is_train:
assert (max_class == len(self.classes) - 1)
self.lb_ids = [self.dict[el] for el in self.imgs]
self.lb_cams = [-1 for el in self.imgs]
self.imgs = [os.path.join(data_path, el) for el in self.imgs]
if is_train:
self.trans = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.RandomCrop((img_size, img_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.486, 0.459, 0.408),
(0.229, 0.224, 0.225)),
RandomErasing(0.5, mean=[0.0, 0.0, 0.0])
])
else:
self.trans_tuple = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.486, 0.459, 0.408),
(0.229, 0.224, 0.225))
])
self.Lambda = transforms.Lambda(
lambda crops: [self.trans_tuple(crop) for crop in crops])
self.trans = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.TenCrop((img_size, img_size)),
self.Lambda,
])
# useful for sampler
self.lb_img_dict = dict()
self.lb_ids_uniq = set(self.lb_ids)
lb_array = np.array(self.lb_ids)
for lb in self.lb_ids_uniq:
idx = np.where(lb_array == lb)[0]
self.lb_img_dict.update({lb: idx})
from train_config import args, params
grad_list = []
def print_grad(grad):
print(grad)
grad_list.append(grad)
data_transforms = {
'train': transforms.Compose([
transforms.Resize((224,224)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.1,0.1,0.1,0),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
RandomErasing(probability=0.5, mean=[0.0, 0.0, 0.0])
]),
'val': transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
train_datasets = PETAData(txt_path='data_list/PETA_RAP_train.txt',dataset='train', data_transforms = data_transforms)
test_datasets = PETAData(txt_path='data_list/PETA_RAP_test.txt', dataset='val', data_transforms = data_transforms)
image_datasets = {'train' : train_datasets, 'val' : test_datasets}
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))
def __init__(self, **kwargs):
super(RandomErasingImageAugmentor, self).__init__(**kwargs)
additional_transforms = [RandomErasing()]
self.add_transforms(additional_transforms)
def __init__(self, hyperparameters, gpu_ids=[0]):
super(DGNet_Trainer, self).__init__()
# 从配置文件获取生成模型和鉴别模型的学习率
lr_g = hyperparameters['lr_g']
lr_d = hyperparameters['lr_d']
# ID 类别
ID_class = hyperparameters['ID_class']
# 是否设置使用fp16，
if not 'apex' in hyperparameters.keys():
hyperparameters['apex'] = False
self.fp16 = hyperparameters['apex']
# Initiate the networks
# We do not need to manually set fp16 in the network for the new apex. So here I set fp16=False.
self.gen_a = AdaINGen(hyperparameters['input_dim_a'],
hyperparameters['gen'],
fp16=False) # auto-encoder for domain a
self.gen_b = self.gen_a # auto-encoder for domain b
'''
ft_netAB : Ea
'''
# ID_stride： 外观编码器池化层的stride
if not 'ID_stride' in hyperparameters.keys():
hyperparameters['ID_stride'] = 2
# id_a ： 外观编码器 -> Ea
if hyperparameters['ID_style'] == 'PCB':
self.id_a = PCB(ID_class)
elif hyperparameters['ID_style'] == 'AB':
self.id_a = ft_netAB(ID_class,
stride=hyperparameters['ID_stride'],
norm=hyperparameters['norm_id'],
pool=hyperparameters['pool'])
else:
self.id_a = ft_net(ID_class,
norm=hyperparameters['norm_id'],
pool=hyperparameters['pool']) # return 2048 now
self.id_b = self.id_a # 对图片b的操作与图片a的操作一致
# 判别器，使用的是一个多尺寸的判别器，就是对图片进行几次缩放，并且对每次缩放都会预测，计算总的损失
# 经过网络3个缩放，，分别为：[batch_size, 1, 64, 32],[batch_size, 1, 32, 16],[batch_size, 1, 16, 8]
self.dis_a = MsImageDis(3, hyperparameters['dis'],
fp16=False) # discriminator for domain a
self.dis_b = self.dis_a # discriminator for domain b
# load teachers
if hyperparameters['teacher'] != "":
teacher_name = hyperparameters['teacher']
print(teacher_name)
# 加载多个老师模型
teacher_names = teacher_name.split(',')
# 构建老师模型
teacher_model = nn.ModuleList() # 初始化为空，接下来开始填充
teacher_count = 0
for teacher_name in teacher_names:
config_tmp = load_config(teacher_name)
# 池化层的stride
if 'stride' in config_tmp:
stride = config_tmp['stride']
else:
stride = 2
# 开始搭建网络
model_tmp = ft_net(ID_class, stride=stride)
teacher_model_tmp = load_network(model_tmp, teacher_name)
teacher_model_tmp.model.fc = nn.Sequential(
) # remove the original fc layer in ImageNet
teacher_model_tmp = teacher_model_tmp.cuda()
# teacher_model_tmp,[3, 224, 224]
# 使用fp16
if self.fp16:
teacher_model_tmp = amp.initialize(teacher_model_tmp,
opt_level="O1")
teacher_model.append(teacher_model_tmp.cuda().eval(
)) # 第一个填充为 teacher_model_tmp.cuda().eval()
teacher_count += 1
self.teacher_model = teacher_model
# 是否使用batchnorm
if hyperparameters['train_bn']:
self.teacher_model = self.teacher_model.apply(train_bn)
# 实例正则化
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# RGB to one channel
# 因为Es 需要使用灰度图， 所以single 用来将图片转化为灰度图
if hyperparameters['single'] == 'edge':
self.single = to_edge
else:
self.single = to_gray(False)
# Random Erasing when training
# arasing_p 随机擦除的概率
if not 'erasing_p' in hyperparameters.keys():
self.erasing_p = 0
else:
self.erasing_p = hyperparameters['erasing_p']
# 对图片中的某一随机区域进行擦除，具体：将该区域的像素值设置为均值
self.single_re = RandomErasing(probability=self.erasing_p,
mean=[0.0, 0.0, 0.0])
if not 'T_w' in hyperparameters.keys():
hyperparameters['T_w'] = 1
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(
self.dis_a.parameters()) #+ list(self.dis_b.parameters())
gen_params = list(
self.gen_a.parameters()) #+ list(self.gen_b.parameters())
self.dis_opt = torch.optim.Adam(
[p for p in dis_params if p.requires_grad],
lr=lr_d,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr_g,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
# id params
# 修改 id_a模型中分类器的学习率
if hyperparameters['ID_style'] == 'PCB':
ignored_params = (
list(map(id, self.id_a.classifier0.parameters())) +
list(map(id, self.id_a.classifier1.parameters())) +
list(map(id, self.id_a.classifier2.parameters())) +
list(map(id, self.id_a.classifier3.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD(
[{
'params': base_params,
'lr': lr2
}, {
'params': self.id_a.classifier0.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier1.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier2.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier3.parameters(),
'lr': lr2 * 10
}],
weight_decay=hyperparameters['weight_decay'],
momentum=0.9,
nesterov=True)
elif hyperparameters['ID_style'] == 'AB':
ignored_params = (
list(map(id, self.id_a.classifier1.parameters())) +
list(map(id, self.id_a.classifier2.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD(
[{
'params': base_params,
'lr': lr2
}, {
'params': self.id_a.classifier1.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier2.parameters(),
'lr': lr2 * 10
}],
weight_decay=hyperparameters['weight_decay'],
momentum=0.9,
nesterov=True)
else:
ignored_params = list(map(id, self.id_a.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD(
[{
'params': base_params,
'lr': lr2
}, {
'params': self.id_a.classifier.parameters(),
'lr': lr2 * 10
}],
weight_decay=hyperparameters['weight_decay'],
momentum=0.9,
nesterov=True)
# 生成器和判别器中的优化策略（学习率的更新策略）
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
self.id_scheduler = get_scheduler(self.id_opt, hyperparameters)
self.id_scheduler.gamma = hyperparameters['gamma2']
#ID Loss
self.id_criterion = nn.CrossEntropyLoss()
self.criterion_teacher = nn.KLDivLoss(
size_average=False) # 生成主要特征： Lprim
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] +
'/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
# save memory
# 保存当前的模型，是为了提高计算效率
if self.fp16:
# Name the FP16_Optimizer instance to replace the existing optimizer
assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled."
self.gen_a = self.gen_a.cuda()
self.dis_a = self.dis_a.cuda()
self.id_a = self.id_a.cuda()
self.gen_b = self.gen_a
self.dis_b = self.dis_a
self.id_b = self.id_a
self.gen_a, self.gen_opt = amp.initialize(self.gen_a,
self.gen_opt,
opt_level="O1")
self.dis_a, self.dis_opt = amp.initialize(self.dis_a,
self.dis_opt,
opt_level="O1")
self.id_a, self.id_opt = amp.initialize(self.id_a,
self.id_opt,
opt_level="O1")
def __init__(self, hyperparameters, gpu_ids=[0]):
super(DGNet_Trainer, self).__init__()
# 从配置文件获取生成模型的和鉴别模型的学习率
lr_g = hyperparameters['lr_g']
lr_d = hyperparameters['lr_d']
# # ID的类别，这里要注意，不同的数据集都是不一样的，应该是训练数据集的ID数目，非测试集
ID_class = hyperparameters['ID_class']
# 看是否设置使用float16，估计float16可以增加精确度
if not 'apex' in hyperparameters.keys():
hyperparameters['apex'] = False
self.fp16 = hyperparameters['apex']
# Initiate the networks
# We do not need to manually set fp16 in the network for the new apex. So here I set fp16=False.
################################################################################################################
##这里是定义Es和G
# 注意这里包含了两个步骤，Es编码+解码过程，既然解码（论文Figure 2的黄色梯形G）包含到这里了，下面Ea应该不会包含解码过程了
# 因为这里是一个类，如后续gen_a.encode()可以进行编码，gen_b.encode()可以进行解码
self.gen_a = AdaINGen(hyperparameters['input_dim_a'],
hyperparameters['gen'],
fp16=False) # auto-encoder for domain a
self.gen_b = self.gen_a # auto-encoder for domain b
############################################################################################################################################
############################################################################################################################################
##这里是定义Ea
# ID_stride，外观编码器池化层的stride
if not 'ID_stride' in hyperparameters.keys():
hyperparameters['ID_stride'] = 2
# hyperparameters['ID_style']默认为'AB'，论文中的Ea编码器
#这里是设置Ea，有三种模型可以选择
#PCB模型，ft_netAB为改造后的resnet50，ft_net为resnet50
if hyperparameters['ID_style'] == 'PCB':
self.id_a = PCB(ID_class)
elif hyperparameters['ID_style'] == 'AB':
# 这是我们执行的模型，注意的是，id_a返回两个x（表示身份），获得f，具体介绍看函数内部
# 我们使用的是ft_netAB，是代码中Ea编码的过程，也就得到 ap code的过程，除了ap code还会得到两个分类结果
# 现在怀疑，该分类结果，可能就是行人重识别的结果
#ID_class表示有ID_class个不同ID的行人
self.id_a = ft_netAB(ID_class,
stride=hyperparameters['ID_stride'],
norm=hyperparameters['norm_id'],
pool=hyperparameters['pool'])
else:
self.id_a = ft_net(ID_class,
norm=hyperparameters['norm_id'],
pool=hyperparameters['pool']) # return 2048 now
# 这里进行的是浅拷贝，所以我认为他们的权重是一起的，可以理解为一个
self.id_b = self.id_a
############################################################################################################################################################
############################################################################################################################################################
##这里是定义D
# 鉴别器，行人重识别，这里使用的是一个多尺寸的鉴别器，大概就是说，对图片进行几次缩放,并且对每次缩放都会预测，计算总的损失
# 经过网络3个元素，分别大小为[batch_size,1，64，32], [batch_size,1，32，16], [batch_size,1，16，8]
self.dis_a = MsImageDis(3, hyperparameters['dis'],
fp16=False) # discriminator for domain a
self.dis_b = self.dis_a # discriminator for domain b
############################################################################################################################################################
############################################################################################################################################################
# load teachers
# 加载老师模型
# teacher:老师模型名称。对于DukeMTMC，您可以设置“best - duke”
if hyperparameters['teacher'] != "":
#teacher_name=best
teacher_name = hyperparameters['teacher']
print(teacher_name)
#有这个操作，我怀疑是可以加载多个教师模型
teacher_names = teacher_name.split(',')
#构建老师模型
teacher_model = nn.ModuleList()
teacher_count = 0
# 默认只有一个teacher_name='teacher_name'，所以其加载的模型配置文件为项目根目录models/best/opts.yaml模型
for teacher_name in teacher_names:
# 加载配置文件models/best/opts.yaml
config_tmp = load_config(teacher_name)
if 'stride' in config_tmp:
#stride=1
stride = config_tmp['stride']
else:
stride = 2
# 老师模型加载，老师模型为ft_net为resnet50
model_tmp = ft_net(ID_class, stride=stride)
teacher_model_tmp = load_network(model_tmp, teacher_name)
# 移除原本的全连接层
teacher_model_tmp.model.fc = nn.Sequential(
) # remove the original fc layer in ImageNet
teacher_model_tmp = teacher_model_tmp.cuda()
# summary(teacher_model_tmp, (3, 224, 224))
#使用浮点型
if self.fp16:
teacher_model_tmp = amp.initialize(teacher_model_tmp,
opt_level="O1")
teacher_model.append(teacher_model_tmp.cuda().eval())
teacher_count += 1
self.teacher_model = teacher_model
# 选择是否使用bn
if hyperparameters['train_bn']:
self.teacher_model = self.teacher_model.apply(train_bn)
############################################################################################################################################################
# 实例正则化
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# RGB to one channel
# 默认设置signal=gray，Es的输入为灰度图
if hyperparameters['single'] == 'edge':
self.single = to_edge
else:
self.single = to_gray(False)
# Random Erasing when training
#earsing_p表示随机擦除的概率
if not 'erasing_p' in hyperparameters.keys():
self.erasing_p = 0
else:
self.erasing_p = hyperparameters['erasing_p']
#随机擦除矩形区域的一些像素，应该类似于数据增强
self.single_re = RandomErasing(probability=self.erasing_p,
mean=[0.0, 0.0, 0.0])
# 设置T_w为1,T_w为primary feature learning loss的权重系数
if not 'T_w' in hyperparameters.keys():
hyperparameters['T_w'] = 1
################################################################################################
# Setup the optimizers
# 设置优化器参数
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(
self.dis_a.parameters()) #+ list(self.dis_b.parameters())
gen_params = list(
self.gen_a.parameters()) #+ list(self.gen_b.parameters())
#使用Adams优化器，用Adams训练Es，G,D
self.dis_opt = torch.optim.Adam(
[p for p in dis_params if p.requires_grad],
lr=lr_d,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr_g,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
# id params
# 因为ID_style默认为AB，所以这里不执行
if hyperparameters['ID_style'] == 'PCB':
ignored_params = (
list(map(id, self.id_a.classifier0.parameters())) +
list(map(id, self.id_a.classifier1.parameters())) +
list(map(id, self.id_a.classifier2.parameters())) +
list(map(id, self.id_a.classifier3.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
self.id_a.parameters())
lr2 = hyperparameters['lr2']
#Ea 的优化器
self.id_opt = torch.optim.SGD(
[{
'params': base_params,
'lr': lr2
}, {
'params': self.id_a.classifier0.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier1.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier2.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier3.parameters(),
'lr': lr2 * 10
}],
weight_decay=hyperparameters['weight_decay'],
momentum=0.9,
nesterov=True)
# 这里是我们执行的代码
elif hyperparameters['ID_style'] == 'AB':
# 忽略的参数，应该是适用于'PCB'或者其他的，但是不适用于'AB'的
ignored_params = (
list(map(id, self.id_a.classifier1.parameters())) +
list(map(id, self.id_a.classifier2.parameters())))
# 获得基本的配置参数，如学习率
base_params = filter(lambda p: id(p) not in ignored_params,
self.id_a.parameters())
lr2 = hyperparameters['lr2']
#对Ea使用SGD
self.id_opt = torch.optim.SGD(
[{
'params': base_params,
'lr': lr2
}, {
'params': self.id_a.classifier1.parameters(),
'lr': lr2 * 10
}, {
'params': self.id_a.classifier2.parameters(),
'lr': lr2 * 10
}],
weight_decay=hyperparameters['weight_decay'],
momentum=0.9,
nesterov=True)
else:
ignored_params = list(map(id, self.id_a.classifier.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD(
[{
'params': base_params,
'lr': lr2
}, {
'params': self.id_a.classifier.parameters(),
'lr': lr2 * 10
}],
weight_decay=hyperparameters['weight_decay'],
momentum=0.9,
nesterov=True)
# 选择各个网络的优化
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
self.id_scheduler = get_scheduler(self.id_opt, hyperparameters)
self.id_scheduler.gamma = hyperparameters['gamma2']
#ID Loss
#交叉熵损失函数
self.id_criterion = nn.CrossEntropyLoss()
# KL散度
self.criterion_teacher = nn.KLDivLoss(size_average=False)
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] +
'/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
# save memory
if self.fp16:
# Name the FP16_Optimizer instance to replace the existing optimizer
assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled."
self.gen_a = self.gen_a.cuda()
self.dis_a = self.dis_a.cuda()
self.id_a = self.id_a.cuda()
self.gen_b = self.gen_a
self.dis_b = self.dis_a
self.id_b = self.id_a
self.gen_a, self.gen_opt = amp.initialize(self.gen_a,
self.gen_opt,
opt_level="O1")
self.dis_a, self.dis_opt = amp.initialize(self.dis_a,
self.dis_opt,
opt_level="O1")
self.id_a, self.id_opt = amp.initialize(self.id_a,
self.id_opt,
opt_level="O1")
def run(args):
gpuId, epochs, weight_decay, batch_id, batch_image, lr_1, lr_2, erasing_p, sampling, exp_dir, trainset_name, cd_trainset_name, testset_names, rand_crop, head_1part_stride = \
args.gpuId, args.epochs, args.weight_decay, args.batch_id, args.batch_image, args.lr_1, args.lr_2, args.erasing_p, args.sampling, args.exp_dir, args.trainset_name, args.cd_trainset_name, args.testset_names, args.rand_crop, args.head_1part_stride
DEVICE = torch.device("cuda:" +
gpuId if torch.cuda.is_available() else "cpu")
print(DEVICE)
num_workers = 4
batch_test = 64 #32
train_list = [
transforms.Resize((400, 144)),
transforms.RandomCrop((384, 128))
] if rand_crop else [transforms.Resize((384, 128))]
train_list += [
transforms.ToTensor(),
]
if erasing_p > 0:
train_list = train_list + [
RandomErasing(probability=erasing_p, mean=[0.0, 0.0, 0.0])
]
train_list += [
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
train_transform = transforms.Compose(train_list)
if trainset_name in ['market1501', 'cuhk03', 'duke']:
root = get_dataset_root(trainset_name)
train_dataset = Market1501(
root + '/bounding_box_train',
transform=train_transform,
training=True,
kpt_file=trainset_name + '-kpt.pkl' if args.pap else None,
ps_dir=root + '_ps_label' if args.src_ps_lw > 0 else None)
elif trainset_name in ['msmt17']:
train_dataset = MSMT17(transform=train_transform,
training=True,
use_kpt=args.pap,
use_ps=args.src_ps_lw > 0,
split='train')
else:
raise ValueError('Invalid train set {}'.format(trainset_name))
train_loader = DataLoader(train_dataset,
sampler=RandomIdSampler(train_dataset,
batch_image=batch_image),
batch_size=batch_id * batch_image,
num_workers=num_workers,
drop_last=True)
if args.cd_ps_lw > 0:
if cd_trainset_name in ['market1501', 'cuhk03', 'duke']:
cd_train_dataset = Market1501(get_dataset_root(cd_trainset_name) +
'/bounding_box_train',
transform=train_transform,
training=True,
kpt_file=None,
ps_dir=cd_trainset_name + '-ps')
elif cd_trainset_name in ['msmt17']:
cd_train_dataset = MSMT17(transform=train_transform,
training=True,
use_kpt=False,
use_ps=True)
else:
raise ValueError(
'Invalid cd train set {}'.format(cd_trainset_name))
cd_train_loader = InfiniteNextBatch(
DataLoader(cd_train_dataset,
batch_size=args.cd_train_batch_size,
num_workers=num_workers,
drop_last=True))
test_transform = transforms.Compose([
transforms.Resize((384, 128)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
test_flip_transform = transforms.Compose([
transforms.Resize((384, 128)), functional.hflip,
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
def make_test_loader_M_C_D(root, name):
query_dataset = Market1501(root + '/query',
transform=test_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
query_flip_dataset = Market1501(root + '/query',
transform=test_flip_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
query_loader = DataLoader(query_dataset,
batch_size=batch_test,
num_workers=num_workers)
query_flip_loader = DataLoader(query_flip_dataset,
batch_size=batch_test,
num_workers=num_workers)
test_dataset = Market1501(root + '/bounding_box_test',
transform=test_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
test_flip_dataset = Market1501(root + '/bounding_box_test',
transform=test_flip_transform,
training=False,
kpt_file=name +
'-kpt.pkl' if args.pap else None)
test_loader = DataLoader(test_dataset,
batch_size=batch_test,
num_workers=num_workers)
test_flip_loader = DataLoader(test_flip_dataset,
batch_size=batch_test,
num_workers=num_workers)
return query_loader, query_flip_loader, test_loader, test_flip_loader
def make_test_loader_MS_PR_PI(name):
if name == 'msmt17':
dclass = MSMT17
elif name == 'partial_reid':
dclass = PartialREID
elif name == 'partial_ilids':
dclass = PartialiLIDs
else:
raise ValueError('Invalid dataset name {}'.format(name))
q_set = dclass(transform=test_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='query')
q_flip_set = dclass(transform=test_flip_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='query')
q_loader = DataLoader(q_set,
batch_size=batch_test,
num_workers=num_workers)
q_flip_loader = DataLoader(q_flip_set,
batch_size=batch_test,
num_workers=num_workers)
g_set = dclass(transform=test_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='gallery')
g_flip_set = dclass(transform=test_flip_transform,
training=False,
use_kpt=args.pap,
use_ps=False,
split='gallery')
g_loader = DataLoader(g_set,
batch_size=batch_test,
num_workers=num_workers)
g_flip_loader = DataLoader(g_flip_set,
batch_size=batch_test,
num_workers=num_workers)
return q_loader, q_flip_loader, g_loader, g_flip_loader
def make_test_loader(name):
if name in ['market1501', 'cuhk03', 'duke']:
return make_test_loader_M_C_D(get_dataset_root(name), name)
elif name in ['msmt17', 'partial_reid', 'partial_ilids']:
return make_test_loader_MS_PR_PI(name)
test_loaders = [make_test_loader(name) for name in testset_names]
mgn = MGN(len(train_dataset.unique_ids), args)
if torch.cuda.device_count() > 1:
mgn = nn.DataParallel(mgn)
mgn = mgn.to(DEVICE)
vanilla_cross_entropy_loss = nn.CrossEntropyLoss()
cross_entropy_loss = nn.CrossEntropyLoss(reduce=False)
triplet_semihard_loss = TripletSemihardLoss(
margin=0.1,
DEVICE=DEVICE,
sampling=sampling,
batch_id=batch_id,
batch_image=batch_image) #batch_hard, .'curriculum'
ps_loss = PSLoss()
optimizer_start1 = optim.SGD(mgn.parameters(),
lr=lr_1,
momentum=0.9,
weight_decay=weight_decay)
optimizer_start2 = optim.SGD(mgn.parameters(),
lr=lr_2,
momentum=0.9,
weight_decay=weight_decay)
scheduler_1 = optim.lr_scheduler.MultiStepLR(optimizer_start1, [140, 180],
gamma=0.1)
scheduler_2 = optim.lr_scheduler.MultiStepLR(
optimizer_start2, [140, 180], gamma=0.1) # best [140, 180] [120, 160]
def get_model_input(inputs, target):
dic = {'im': inputs.to(DEVICE)}
if 'pap_mask_2p' in target:
dic['pap_mask_2p'] = target['pap_mask_2p'].to(DEVICE)
dic['pap_mask_3p'] = target['pap_mask_3p'].to(DEVICE)
return dic
def extract_loader_feat(loader, verbose=False):
feat = []
vis = []
i = 0
for inputs, target in loader:
if verbose:
print(i)
i += 1
with torch.no_grad():
output = mgn(get_model_input(inputs, target))
feat.append(output[1].detach().cpu().numpy())
if args.pap:
vis_ = np.concatenate([
np.ones([len(output[1]), 3]),
torch.stack(output[5 + 3 + 5:5 + 3 + 5 + 5],
1).detach().cpu().numpy()
], 1)
vis.append(vis_)
feat = np.concatenate(feat)
vis = np.concatenate(vis) if args.pap else None
return feat, vis
def test(query_loader,
query_flip_loader,
test_loader,
test_flip_loader,
trainset_name,
testset_name,
epoch,
verbose=False):
cache_file = '{}/feat_cache-{}_to_{}.pkl'.format(
exp_dir, trainset_name, testset_name)
if args.use_feat_cache:
assert os.path.exists(
cache_file), "Feature cache file {} does not exist!".format(
cache_file)
query_2, q_vis, query_flip_2, q_vis, test_2, test_vis, test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams = load_pickle(
cache_file)
else:
query_2, q_vis = extract_loader_feat(query_loader, verbose=verbose)
query_flip_2, q_vis = extract_loader_feat(query_flip_loader,
verbose=verbose)
test_2, test_vis = extract_loader_feat(test_loader,
verbose=verbose)
test_flip_2, test_vis = extract_loader_feat(test_flip_loader,
verbose=verbose)
q_ids = query_loader.dataset.ids
q_cams = query_loader.dataset.cameras
g_ids = test_loader.dataset.ids
g_cams = test_loader.dataset.cameras
save_pickle([
query_2, q_vis, query_flip_2, q_vis, test_2, test_vis,
test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams
], cache_file)
if args.test_which_feat > 0:
# TODO: implement for pap
idx = args.test_which_feat
query_2 = query_2[:, 256 * idx - 256:256 * idx]
query_flip_2 = query_flip_2[:, 256 * idx - 256:256 * idx]
test_2 = test_2[:, 256 * idx - 256:256 * idx]
test_flip_2 = test_flip_2[:, 256 * idx - 256:256 * idx]
query = normalize(query_2 + query_flip_2)
test = normalize(test_2 + test_flip_2)
if verbose:
print('query.shape:', query.shape)
print('test.shape:', test.shape)
if args.pap:
print('q_vis.shape:', q_vis.shape)
print('test_vis.shape:', test_vis.shape)
if args.pap:
dist_1 = compute_dist_with_visibility(query,
test,
q_vis,
test_vis,
dist_type='euclidean',
avg_by_vis_num=False)
else:
dist_1 = cdist(query, test)
r_1 = cmc(dist_1,
q_ids,
g_ids,
q_cams,
g_cams,
separate_camera_set=False,
single_gallery_shot=False,
first_match_break=True)
m_ap_1 = mean_ap(dist_1, q_ids, g_ids, q_cams, g_cams)
print('EPOCH [%d] %s -> %s: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' %
(epoch + 1, trainset_name, testset_name, m_ap_1, r_1[0], r_1[2],
r_1[4], r_1[9]))
if args.only_test:
mgn.eval()
if not args.use_feat_cache:
if args.model_weight_file:
model_weight_file = args.model_weight_file
else:
model_weight_file = '{}/model_weight.pth'.format(exp_dir)
load_model_weight((mgn.module if hasattr(mgn, 'module') else mgn),
model_weight_file)
for name, test_loader in zip(testset_names, test_loaders):
test(test_loader[0],
test_loader[1],
test_loader[2],
test_loader[3],
trainset_name,
name,
-1,
verbose=False)
exit()
for epoch in range(epochs):
mgn.train()
scheduler_1.step()
scheduler_2.step()
running_loss = 0.0
running_loss_1 = 0.0
running_loss_2 = 0.0
if epoch < 20:
optimizer_1 = optim.SGD(mgn.parameters(),
lr=0.01 + 0.0045 * epoch,
momentum=0.9,
weight_decay=weight_decay)
optimizer_2 = optim.SGD(mgn.parameters(),
lr=0.001 + 0.00045 * epoch,
momentum=0.9,
weight_decay=weight_decay)
else:
optimizer_1 = optimizer_start1
optimizer_2 = optimizer_start2
for i, data in enumerate(train_loader):
inputs, target = data
inputs = inputs.to(DEVICE)
for k, v in target.items():
target[k] = v.to(DEVICE)
labels = target['id']
outputs = mgn(get_model_input(inputs, target))
optimizer_1.zero_grad()
if args.pap:
losses_1 = [
vanilla_cross_entropy_loss(output, labels)
for output in outputs[5:5 + 3]
] + [(cross_entropy_loss(output, labels) * v).sum() /
(v.sum() + 1e-12)
for output, v in zip(outputs[5 + 3:5 + 3 +
5], outputs[5 + 3 + 5:5 + 3 +
5 + 5])]
else:
losses_1 = [
vanilla_cross_entropy_loss(output, labels)
for output in outputs[5:5 + 8]
]
loss_1 = sum(losses_1) / len(losses_1)
psl = 0
if args.src_ps_lw > 0:
psl = (ps_loss(outputs[-3], target['ps_label']) +
ps_loss(outputs[-2], target['ps_label']) +
ps_loss(outputs[-1], target['ps_label'])) / 3.
(loss_1 + psl * args.src_ps_lw).backward()
if args.cd_ps_lw > 0:
cd_inputs, cd_targets = cd_train_loader.next_batch()
cd_inputs = cd_inputs.to(DEVICE)
for k, v in cd_targets.items():
cd_targets[k] = v.to(DEVICE)
pap_old = args.pap
args.pap = False
outputs = mgn(get_model_input(cd_inputs, cd_targets))
args.pap = pap_old
cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) +
ps_loss(outputs[-2], cd_targets['ps_label']) +
ps_loss(outputs[-1], cd_targets['ps_label'])) / 3.
(cd_psl * args.cd_ps_lw).backward()
optimizer_1.step()
outputs = mgn(get_model_input(inputs, target))
optimizer_2.zero_grad()
losses_2 = [
triplet_semihard_loss(output, labels, epoch)
for output in outputs[2:5]
]
loss_2 = sum(losses_2) / len(losses_2)
psl = 0
if args.src_ps_lw > 0:
psl = (ps_loss(outputs[-3], target['ps_label']) +
ps_loss(outputs[-2], target['ps_label']) +
ps_loss(outputs[-1], target['ps_label'])) / 3.
(loss_2 + psl * args.src_ps_lw).backward()
if args.cd_ps_lw > 0:
cd_inputs, cd_targets = cd_train_loader.next_batch()
cd_inputs = cd_inputs.to(DEVICE)
for k, v in cd_targets.items():
cd_targets[k] = v.to(DEVICE)
pap_old = args.pap
args.pap = False
outputs = mgn(get_model_input(cd_inputs, cd_targets))
args.pap = pap_old
cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) +
ps_loss(outputs[-2], cd_targets['ps_label']) +
ps_loss(outputs[-1], cd_targets['ps_label'])) / 3.
(cd_psl * args.cd_ps_lw).backward()
optimizer_2.step()
running_loss_1 += loss_1.item()
running_loss_2 += loss_2.item()
running_loss = running_loss + (loss_1.item() + loss_2.item()) / 2.0
print('%d/%d - %d/%d - loss: %f - ps_loss: %f - cd_ps_loss: %f' %
(epoch + 1, epochs, i, len(train_loader),
(loss_1.item() + loss_2.item()) / 2,
psl.item() if isinstance(psl, torch.Tensor) else 0,
cd_psl.item() if args.cd_ps_lw > 0 else 0))
print('epoch: %d/%d - loss1: %f' %
(epoch + 1, epochs, running_loss_1 / len(train_loader)))
print('epoch: %d/%d - loss2: %f' %
(epoch + 1, epochs, running_loss_2 / len(train_loader)))
# if (epoch + 1) % 50 == 0:
# model_weight_file = '{}/model_weight.pth'.format(exp_dir)
# save_model(mgn, model_weight_file)
# mgn.eval()
# for name, test_loader in zip(testset_names, test_loaders):
# test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, epoch)
model_weight_file = '{}/model_weight.pth'.format(exp_dir)
save_model(mgn, model_weight_file)
mgn.eval()
for name, test_loader in zip(testset_names, test_loaders):
test(test_loader[0], test_loader[1], test_loader[2], test_loader[3],
trainset_name, name, epoch)
# Load Data
transform_train_list = [
transforms.Resize((384, 192), interpolation=transforms.InterpolationMode.BICUBIC),
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=transforms.InterpolationMode.BICUBIC),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if opt.erase_prob > 0:
transform_train_list += [RandomErasing(probability=opt.erase_prob, 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)]
print('dataset transformation: ', transform_train_list)
data_transforms = {
'train': transforms.Compose(transform_train_list),
'val': transforms.Compose(transform_val_list),
}
image_datasets = {
key: datasets.ImageFolder(os.path.join(data_dir, key), value)
for key, value in data_transforms.items()
}
random_height_crop = 224
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.Resize(256,interpolation=3),
# transforms.RandomCrop(224,224),
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(opt.erasing_p)
]
# print(transform_train_list)
transform_val_list = [
transforms.Resize(size=(256, 128), interpolation=3), # Image.BICUBIC
# transforms.Resize(256,interpolation=3),
# transforms.RandomCrop(224,224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
data_transforms = {
'train': transforms.Compose(transform_train_list),
'val': transforms.Compose(transform_val_list),
transform_train_list = [
transforms.Resize((384, 128), 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, 128), 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 = ''
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
# tensorboardX
# writer = SummaryWriter(log_dir=osp.join(args.save_dir,'summary'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_img_dataset(
root=args.root, name=args.dataset, split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split,
)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
if args.random_erasing:
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
RandomErasing(probability=args.probability, mean=[0.0, 0.0, 0.0]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
if args.loss == 'xent,htri':
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances),
batch_size=args.train_batch, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
elif args.loss == 'xent':
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
batch_size=args.train_batch, shuffle=True, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss=args.loss)
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0))
criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
if args.stepsize > 0:
if not args.warmup:
scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)
start_epoch = args.start_epoch
if args.resume:
print("Loading checkpoint from '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test(model, queryloader, galleryloader, use_gpu)
return
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
#lr = 1e-3 * len(args.gpu_devices)
lr = 1e-3
elif ep < 180:
#lr = 1e-4 * len(args.gpu_devices)
lr = 1e-4
elif ep < 300:
#lr = 1e-5 * len(args.gpu_devices)
lr = 1e-5
elif ep < 320:
#lr = 1e-5 * 0.1 ** ((ep - 320) / 80) * len(args.gpu_devices)
lr = 1e-5 * 0.1 ** ((ep - 320) / 80)
elif ep < 400:
lr = 1e-6
elif ep < 480:
#lr = 1e-4 * len(args.gpu_devices)
lr = 1e-4
else:
#lr = 1e-5 * len(args.gpu_devices)
lr = 1e-5
for p in optimizer.param_groups:
p['lr'] = lr
length = len(trainloader)
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
#best_rerank1 = -np.inf
#best_rerankepoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
if args.stepsize > 0:
if args.warmup:
adjust_lr(optimizer, epoch + 1)
else:
scheduler.step()
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu=use_gpu, summary=None, length=length)
train_time += round(time.time() - start_train_time)
if (epoch+1) > args.start_eval and args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch:
print("==> Test")
rank1 = test(epoch, model, queryloader, galleryloader, use_gpu=True, summary=None)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
####### Best Rerank
#is_rerankbest = rerank1 > best_rerank1
#if is_rerankbest:
# best_rerank1 = rerank1
# best_rerankepoch = 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(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar'))
writer.close()
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch))
#print("==> Best Rerank-1 {:.1%}, achieved at epoch {}".format(best_rerank1, best_rerankepoch))
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))
def __init__(self, hyperparameters, gpu_ids=[0]):
super(DGNet_Trainer, self).__init__()
lr_g = hyperparameters['lr_g'] #生成器学习率
lr_d = hyperparameters['lr_d'] #判别器学习率
ID_class = hyperparameters['ID_class']
if not 'apex' in hyperparameters.keys():
hyperparameters['apex'] = False
self.fp16 = hyperparameters['apex']
# Initiate the networks
# We do not need to manually set fp16 in the network for the new apex. So here I set fp16=False.
# 构建Es编码+解码过程 gen_a.encode()可以进行编码，gen_b.encode()可以进行解码
self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen'], fp16 = False) # auto-encoder for domain a
self.gen_b = self.gen_a # auto-encoder for domain b
# ID_stride，外观编码器池化层的stride
if not 'ID_stride' in hyperparameters.keys():
hyperparameters['ID_stride'] = 2
# 构建外观编码器
if hyperparameters['ID_style']=='PCB':
self.id_a = PCB(ID_class)
elif hyperparameters['ID_style']=='AB': #使用的AB编码器
self.id_a = ft_netAB(ID_class, stride = hyperparameters['ID_stride'], norm=hyperparameters['norm_id'], pool=hyperparameters['pool'])
else:
self.id_a = ft_net(ID_class, norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) # return 2048 now
# 浅拷贝，两者等同
self.id_b = self.id_a
# 鉴别器，使用的是一个多尺寸的鉴别器，即对图片进行几次缩放,并且对每次缩放都会预测，计算总的损失
self.dis_a = MsImageDis(3, hyperparameters['dis'], fp16 = False) # discriminator for domain a
self.dis_b = self.dis_a # discriminator for domain b
# load teachers 加载教师模型
if hyperparameters['teacher'] != "":
teacher_name = hyperparameters['teacher']
print(teacher_name)
# 构建教师模型
teacher_names = teacher_name.split(',')
teacher_model = nn.ModuleList()
teacher_count = 0
for teacher_name in teacher_names:
config_tmp = load_config(teacher_name)
if 'stride' in config_tmp:
stride = config_tmp['stride']
else:
stride = 2
# 网络搭建
model_tmp = ft_net(ID_class, stride = stride)
teacher_model_tmp = load_network(model_tmp, teacher_name)
teacher_model_tmp.model.fc = nn.Sequential() # remove the original fc layer in ImageNet
teacher_model_tmp = teacher_model_tmp.cuda()
if self.fp16:
teacher_model_tmp = amp.initialize(teacher_model_tmp, opt_level="O1")
teacher_model.append(teacher_model_tmp.cuda().eval())
teacher_count +=1
self.teacher_model = teacher_model
# 选择是否使用bn
if hyperparameters['train_bn']:
self.teacher_model = self.teacher_model.apply(train_bn)
# 实例正则化
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
# RGB to one channel
if hyperparameters['single']=='edge':
self.single = to_edge
else:
self.single = to_gray(False)
# Random Erasing when training
if not 'erasing_p' in hyperparameters.keys():
self.erasing_p = 0
else:
self.erasing_p = hyperparameters['erasing_p'] # erasing_p表示随机擦除的概率
# 随机擦除矩形区域的一些像素，数据增强
self.single_re = RandomErasing(probability = self.erasing_p, mean=[0.0, 0.0, 0.0])
if not 'T_w' in hyperparameters.keys():
hyperparameters['T_w'] = 1
# Setup the optimizers 设置优化器的参数
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(self.dis_a.parameters()) #+ list(self.dis_b.parameters())
gen_params = list(self.gen_a.parameters()) #+ list(self.gen_b.parameters())
# 使用Adam优化器
self.dis_opt = torch.optim.Adam([p for p in dis_params if p.requires_grad],
lr=lr_d, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
lr=lr_g, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
# id params
if hyperparameters['ID_style']=='PCB':
ignored_params = (list(map(id, self.id_a.classifier0.parameters() ))
+list(map(id, self.id_a.classifier1.parameters() ))
+list(map(id, self.id_a.classifier2.parameters() ))
+list(map(id, self.id_a.classifier3.parameters() ))
)
base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD([
{'params': base_params, 'lr': lr2},
{'params': self.id_a.classifier0.parameters(), 'lr': lr2*10},
{'params': self.id_a.classifier1.parameters(), 'lr': lr2*10},
{'params': self.id_a.classifier2.parameters(), 'lr': lr2*10},
{'params': self.id_a.classifier3.parameters(), 'lr': lr2*10}
], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True)
elif hyperparameters['ID_style']=='AB':
ignored_params = (list(map(id, self.id_a.classifier1.parameters()))
+ list(map(id, self.id_a.classifier2.parameters())))
# 获得基本的配置参数，如学习率
base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD([
{'params': base_params, 'lr': lr2},
{'params': self.id_a.classifier1.parameters(), 'lr': lr2*10},
{'params': self.id_a.classifier2.parameters(), 'lr': lr2*10}
], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True)
else:
ignored_params = list(map(id, self.id_a.classifier.parameters() ))
base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters())
lr2 = hyperparameters['lr2']
self.id_opt = torch.optim.SGD([
{'params': base_params, 'lr': lr2},
{'params': self.id_a.classifier.parameters(), 'lr': lr2*10}
], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True)
# 选择各个网络优化的策略
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
self.id_scheduler = get_scheduler(self.id_opt, hyperparameters)
self.id_scheduler.gamma = hyperparameters['gamma2']
#ID Loss
self.id_criterion = nn.CrossEntropyLoss()
self.criterion_teacher = nn.KLDivLoss(size_average=False)
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
# save memory
if self.fp16:
# Name the FP16_Optimizer instance to replace the existing optimizer
assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled."
self.gen_a = self.gen_a.cuda()
self.dis_a = self.dis_a.cuda()
self.id_a = self.id_a.cuda()
self.gen_b = self.gen_a
self.dis_b = self.dis_a
self.id_b = self.id_a
self.gen_a, self.gen_opt = amp.initialize(self.gen_a, self.gen_opt, opt_level="O1")
self.dis_a, self.dis_opt = amp.initialize(self.dis_a, self.dis_opt, opt_level="O1")
self.id_a, self.id_opt = amp.initialize(self.id_a, self.id_opt, opt_level="O1")
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 =', transform_train_list)
data_transforms = {
'train': transforms.Compose(transform_train_list),
'val': transforms.Compose(transform_val_list),
}
train_all = ''
# useful for sampler
self.lb_img_dict = dict()
self.lb_ids_uniq = set(self.lb_ids)
lb_array = np.array(self.lb_ids)
for lb in self.lb_ids_uniq:
idx = np.where(lb_array == lb)[0]
self.lb_img_dict.update({lb: idx})
def __len__(self):
return len(self.imgs)
def __getitem__(self, idx):
img = Image.open(self.imgs[idx])
img = self.trans(img)
return img, self.lb_ids[idx], self.lb_cams[idx], self.imgs[idx]
if __name__ == "__main__":
img_dir = "/home/CORP/ryann.bai/dataset/VehicleID/image/"
img_list = "/home/CORP/ryann.bai/dataset/VehicleID/train_test_split_v1/train_list_start0_jpg.txt"
ds = VehicleID(img_dir, img_list, is_train=True)
im, _, _ = ds[1]
print(im.shape)
print(im.max())
print(im.min())
ran_er = RandomErasing()
im = ran_er(im)
cv2.imshow('erased', im)
cv2.waitKey(0)
import os
import numpy as np
from sklearn.model_selection import train_test_split
from runstats import Statistics
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from albumentations import Compose, ShiftScaleRotate, GridDistortion
from albumentations.pytorch import ToTensor
from random_erasing import RandomErasing
random_erasing = RandomErasing()
albumentations_transform = Compose([
ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=0.5),
GridDistortion(),
ToTensor()
])
def get_train_val_data(inp_txt):
with open(inp_txt, 'r') as f:
lines = f.readlines()
lines = [line.rstrip() for line in lines]
fnames = [line.split('\t')[0] for line in lines]
labels = [line.split('\t')[-1] for line in lines]
ref_labels = list(sorted(set(labels)))
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)
train_preprocessing = [
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])
]
val_preprocessing = [
transforms.Resize((256, 128), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]
if ag.erasing_p > 0:
train_preprocessing = train_preprocessing + [RandomErasing(probability=ag.erasing_p, mean=[0.0, 0.0, 0.0])]
if ag.color_jitter:
train_preprocessing = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)]\
+ train_preprocessing
# Compose all the transformations
data_preprocessing = {
'train': transforms.Compose(train_preprocessing),
'val': transforms.Compose(val_preprocessing),
}
# Use all the training samples or not
train_all = ''
if ag.train_all:
train_all = '_all'
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'
args = parser.parse_args()
print('Loading data')
transform_list = [
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])
]
if args.erasing_p > 0:
transform_list.append(
RandomErasing(probability=args.erasing_p, mean=[0.0, 0.0, 0.0]))
if args.color_jitter:
transform_list.append(
transforms.ColorJitter(brightness=0.1,
contrast=0.1,
saturation=0.1,
hue=0))
# Load dataset
transform_list = transforms.Compose(transform_list)
dset_train = MultiViewDataSet(args.data, transform=transform_list)
train_loader = DataLoader(dset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True)
