def __init__(self, args, phase, is_training=True):
self.phase = phase
self.batch_size_ = args.batch_size
# Normalization
norm_value = 255
mean = [
110.63666788 / norm_value, 103.16065604 / norm_value,
96.29023126 / norm_value
]
std = [
38.7568578 / norm_value, 37.88248729 / norm_value,
40.02898126 / norm_value
]
norm_method = transforms.Normalize(mean, std)
# Transforms
if is_training:
assert args.crop_shape is not None and args.crop_shape[
0] == args.crop_shape[1]
crop_method = transforms.MultiScaleRandomCrop(
[1., 0.84, 0.71, 0.59, 0.49], args.crop_shape[0])
spatial_transform = transforms.Compose([
crop_method,
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), norm_method
])
temporal_transform = transforms.TemporalRandomCrop(args.num_slices)
target_transform = transforms.ClassLabel()
n_samples = 1
else:
assert args.crop_shape is not None and args.crop_shape[
0] == args.crop_shape[1]
assert args.resize_shape is not None and args.resize_shape[
0] == args.resize_shape[1]
spatial_transform = transforms.Compose([
transforms.Scale(args.crop_shape[0]),
transforms.CenterCrop(args.crop_shape[0]),
transforms.ToTensor(), norm_method
])
temporal_transform = transforms.LoopPadding(args.num_slices)
target_transform = transforms.ClassLabel()
n_samples = 3
dataset = KineticsDataset(args, phase, n_samples, spatial_transform,
temporal_transform, target_transform)
super(KineticsDataLoader, self).__init__(dataset,
batch_size=args.batch_size,
shuffle=is_training,
num_workers=args.num_workers,
pin_memory=True)
def main():
use_gpu = False
dataset = data_manger.init_img_dataset(root='data', name='market1501',
split_id=False,
cuhk03_labeled=False,
cuhk03_classic_split=False,
)
transforms_test = T.Compose([
T.Resize((256,128)),
# 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]),
])
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transforms_test),
batch_size=30, num_workers=4,
shuffle=False,
pin_memory=False, drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transforms_test),
batch_size=30, num_workers=4,
shuffle=False,
pin_memory=False, drop_last=False,
)
model = models.init_model(name='resnet50', num_classes=751, loss='softmax')
print("Evaluate only")
test(model, queryloader, galleryloader, use_gpu)
def get_coco(root, image_set, transforms, mode='instances'):
anno_file_template = "{}_{}2017.json"
PATHS = {
"train": ("train2017",
os.path.join("annotations",
anno_file_template.format(mode, "train"))),
"val": ("val2017",
os.path.join("annotations",
anno_file_template.format(mode, "val"))),
# "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
}
t = [ConvertCocoPolysToMask()]
if transforms is not None:
t.append(transforms)
transforms = T.Compose(t)
img_folder, ann_file = PATHS[image_set]
img_folder = os.path.join(root, img_folder)
ann_file = os.path.join(root, ann_file)
dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
if image_set == "train":
dataset = _coco_remove_images_without_annotations(dataset)
# dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
return dataset
def single_run(self, pre_load, **kwargs):
""" Performs inference only
Parameters
----------
pre_load : bool
Flag for only loading the model
Returns
-------
The output of the inference
"""
# Load the model if it does not exist yet
if not hasattr(self, 'model') or not hasattr(self, 'transform'):
self.model = self.setup.setup_model(**kwargs)
self.model.eval()
checkpoint = self._load_checkpoint(**kwargs)
if 'test_transform' in checkpoint:
self.transform = checkpoint['test_transform']
else:
logging.info(
"Test transform not found in checkpoint. Using ToTensor()."
)
self.transform = T.Compose([T.ToTensor()])
self.classes = checkpoint['classes']
if pre_load:
# Check no images to process are given
assert kwargs['input_image'] is None
assert kwargs['input_folder'] is None
# Create a fake empty image to process
buffered = io.BytesIO()
Image.new('RGB', (128, 128)).save(buffered, format="PNG")
kwargs['input_image'] = base64.b64encode(
buffered.getvalue()).decode('utf-8')
# Load and preprocess the data
img = self.preprocess(**kwargs)
# Forward Pass
with torch.no_grad():
output = self.model(img)
if pre_load:
# Return a standard answer
payload = {'result': "successfully loaded the model"}
logging.info(f"Returning payload: {payload}")
return payload
else:
# Return post-processed output
return self.postprocess(output, **kwargs)
def main():
# 第四个参数：use_gpu，不需要显示的指定
use_gpu = torch.cuda.is_available()
# if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else False
# 其实可以换一种写法
dataset = data_manager.Market1501(root='data')
# data augmentation
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]),
])
# 第二个参数：queryloader
queryloader = DataLoader(
# 问题：dataset.query哪里来的？ 答：来自dataset = data_manager.Market1501(root='data')
ImageDataset(dataset.query, transform=transform_test),
batch_size=32, shuffle=False, num_workers=4,
pin_memory=pin_memory, drop_last=False,
)
# 第三个参数：galleryloader
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=32, shuffle=False, num_workers=4,
pin_memory=pin_memory, drop_last=False,
)
model = models.init_model(name='resnet50', num_classes=8, loss={'softmax', 'metric'},
aligned=True, use_gpu=use_gpu)
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_class = CrossEntropyLoss(use_gpu=use_gpu)
criterion_metric = TripletLossAlignedReID(margin=0.3)
optimizer = init_optim('adam', model.parameters(), 0.0002, 0.0005)
scheduler = lr_scheduler.StepLR(optimizer, step_size=150, gamma=0.1)
start_epoch = 0
if use_gpu:
model = nn.DataParallel(model).cuda()
# embed()
test(model, queryloader, galleryloader, use_gpu)
return 0
def main():
use_gpu = torch.cuda.is_available()
# use_gpu = False
if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else 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_img_dataset(
root=args.root, name=args.dataset, split_id=args.split_id,
)
print('dataset',dataset)
# data augmentation
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
# T.Resize(size=(384,128),interpolation=3),
T.RandomHorizontalFlip(),
T.ColorJitter(brightness=0.1,contrast=0.1,saturation=0.1,hue=0.1),
# T.RandomVerticalFlip(),
# T.RandomRotation(30),
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.Resize(size=(384,128),interpolation=3),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
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,
)
#embed()
print('len of trainloader',len(trainloader))
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,
)
print('len of queryloader',len(queryloader))
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('len of galleryloader',len(galleryloader))
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_vids,loss={'softmax','metric'},aligned=args.aligned)
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0))
print('Model ',model)
print('num_classes',dataset.num_train_vids)
if args.labelsmooth:
criterion_class = CrossEntropyLabelSmooth(num_classes=dataset.num_train_vids, use_gpu=use_gpu)
else:
# criterion_class = CrossEntropyLoss(use_gpu=use_gpu)
criterion_class = nn.CrossEntropyLoss()
criterion_metric = TripletLossAlignedReID(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
if args.stepsize > 0:
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 0
start_time = time.time()
train_time = 0
best_mAP = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, criterion_metric, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch or ((epoch+1)==1):
print("==> Test")
mAP = test(model, queryloader, galleryloader, use_gpu)
is_best = mAP > best_mAP
if is_best:
best_mAP = mAP
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,
'mAP': mAP,
'epoch': epoch,
}, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print("==> Best mAP {:.2%}, achieved at epoch {}".format(best_mAP, 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 main():
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else 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_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,
)
# data augmentation
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]),
])
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]),
])
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,
)
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={'a_softmax', 'metric'},
aligned=True,
use_gpu=use_gpu)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
if args.labelsmooth:
criterion_class = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
else:
criterion_class = AngleLoss()
criterion_metric = TripletLossAlignedReID(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
if args.stepsize > 0:
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_dict = model.state_dict()
# 1. filter out unnecessary keys
checkpoint = {k: v for k, v in checkpoint.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(checkpoint)
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 0
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, criterion_metric, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
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))
import models
import numpy as np
from util.losses import CrossEntropyLoss, DeepSupervision, CrossEntropyLabelSmooth, TripletLossAlignedReID
from os import getcwd
import torch
import torch.nn as nn
from util.utils import AverageMeter, Logger, save_checkpoint
from custom.testing import test
use_gpu = torch.cuda.is_available()
pin_memory = True if use_gpu else False
transform_test = T.Compose([
T.Resize((256, 128)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
## CUSTOM DATASET
dataset = custom.ValSetCSCE625()
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=4, shuffle=False, num_workers=1,
pin_memory=pin_memory, drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=4, shuffle=False, num_workers=1,
def main():
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if use_gpu:
pin_memory = True
else:
pin_memory = 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))
os.environ['CUDA_CUDA_VISIBLE_DEVICES'] = args.gpu_devices
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
dataset = data_manger.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,
)
#dataloader & augmentation train query gallery
transforms_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]),
])
transforms_test = T.Compose([
T.Resize((args.height, args.width)),
#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]),
])
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transforms_train),
batch_size=args.train_batch,
num_workers=args.workers,
shuffle=True,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transforms_test),
batch_size=args.test_batch,
num_workers=args.workers,
shuffle=False,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transforms_test),
batch_size=args.test_batch,
num_workers=args.workers,
shuffle=False,
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='softmax')
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion = nn.CrossEntropyLoss() #定义损失函数
#optimizer = init_optim(args.optim,model.parameters(),args.lr,args.weight_decay) #定义优化器
# Optimizer
if hasattr(model, 'model'):
base_param_ids = list(map(id, model.model.parameters()))
base_param_ids += list(map(id, model.globe_conv5x.parameters()))
new_params = [
p for p in model.parameters() if id(p) not in base_param_ids
]
param_groups = [{
'params': model.model.parameters(),
'lr_mult': 0.1
}, {
'params': new_params,
'lr_mult': 1.0
}]
else:
param_groups = model.parameters()
optimizer = torch.optim.SGD(param_groups,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay,
nesterov=True)
# ###自己定义优化器
# 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 * args.lr},
# {'params': model.model.fc.parameters(), 'lr': args.lr},
# {'params': model.classifier0.parameters(), 'lr': args.lr},
# {'params': model.classifier1.parameters(), 'lr': args.lr},
# {'params': model.classifier2.parameters(), 'lr': args.lr},
# {'params': model.classifier3.parameters(), 'lr': args.lr},
# # {'params': model.classifier4.parameters(), 'lr': args.lr},
# # {'params': model.classifier5.parameters(), 'lr': args.lr},
# # {'params': model.classifier6.parameters(), 'lr': 0.01},
# # {'params': model.classifier7.parameters(), 'lr': 0.01}
# ], weight_decay=5e-4, momentum=0.9, nesterov=True)
#optimizer = optimizer_ft
# Schedule learning rate
def adjust_lr(epoch):
step_size = 60 if args.arch == 'inception' else args.stepsize
lr = args.lr * (0.1**(epoch // step_size))
for g in optimizer.param_groups:
g['lr'] = lr * g.get('lr_mult', 1)
# if args.stepsize > 0:
# 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_PCB03(model, queryloader, galleryloader, use_gpu)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
adjust_lr(epoch)
start_train_time = time.time()
train_PCB(epoch, model, criterion, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
# if args.stepsize > 0: scheduler.step()
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_PCB02(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 main():
batch_time_total = AverageMeter()
start = time.time()
# 第四个参数：use_gpu，不需要显示的指定
use_gpu = torch.cuda.is_available()
# if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else False
# 其实可以换一种写法
dataset = data_manager.Market1501(root='data')
# data augmentation
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]),
])
# 第二个参数：queryloader
queryloader = DataLoader(
# 问题：dataset.query哪里来的？ 答：来自data_manager中self.query = query
# dataset.query本质为路径集
ImageDataset(dataset.query, transform=transform_test),
batch_size=32,
shuffle=False,
num_workers=4,
pin_memory=pin_memory,
drop_last=False,
)
# 第三个参数：galleryloader
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=32,
shuffle=False,
num_workers=4,
pin_memory=pin_memory,
drop_last=False,
)
model = models.init_model(name='resnet50',
num_classes=8,
loss={'softmax', 'metric'},
aligned=True,
use_gpu=use_gpu)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_class = CrossEntropyLoss(use_gpu=use_gpu)
criterion_metric = TripletLossAlignedReID(margin=0.3)
optimizer = init_optim('adam', model.parameters(), 0.0002, 0.0005)
scheduler = lr_scheduler.StepLR(optimizer, step_size=150, gamma=0.1)
start_epoch = 0
if use_gpu:
model = nn.DataParallel(model).cuda()
# embed()
num, cmc, mAP = test(model, queryloader, galleryloader, use_gpu)
end = time.time()
time_stamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
item_to_json = {
"time_stamp": time_stamp,
"test_results": {
"object_num": num,
"cmc": cmc,
"mAP": mAP,
"time_consumption(s)": end - start
}
}
path = "./output/" + "test_results" + ".json"
s = SaveJson()
s.save_file(path, item_to_json)
# print("==>测试用时: {:.3f} s".format(end - start))
print(" test time(s) | {:.3f}".format(end - start))
print(" ------------------------------")
print("")
# print('------测试结束------')
return 0
def main():
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else 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_img_dataset(
root=args.root,
name=args.dataset,
split_id=args.split_id,
)
# data augmentation
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]),
])
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]),
])
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,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'softmax', 'metric'},
aligned=True,
use_gpu=use_gpu)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
if args.labelsmooth:
criterion_class = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
else:
criterion_class = CrossEntropyLoss(use_gpu=use_gpu)
criterion_metric = TripletLossAlignedReID(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
if args.stepsize > 0:
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()
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
cnt_n = 0
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, criterion_metric, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
cnt_n = cnt_n + 1
if args.stepsize > 0: scheduler.step()
if (cnt_n % 40) == 0: #### Saving models after each 40 epochs
print("==> Saving")
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
is_best = 0
save_checkpoint(
{
'state_dict':
state_dict, ### rank1 and is_best are kept same as original code. Don't want to mess up the saving
'rank1': '###',
'epoch': epoch,
},
is_best,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
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 main():
print(time.strftime("Current TIME is %Y-%m-%d %H:%M:%S", time.localtime()))
torch.manual_seed(args.seed)
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
if use_gpu:
pin_memory = True
else:
pin_memory = False
if not args.evaluate:
# sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) avoid overlay txt file
sys.stdout = Logger(
osp.join(
args.save_dir, "log_train_{}.txt".format(
time.strftime("%Y-%m-%d %H-%M", time.localtime()))))
else:
# sys.stdout = Logger(osp.join(args.save_dir, 'log_test_{}.txt'))
sys.stdout = Logger(
osp.join(
args.save_dir, "log_test_{}.txt".format(
time.strftime("%Y-%m-%d %H-%M", time.localtime()))))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
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)")
# name = 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,
)
# dataloader & augementation train/query/gallery
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]),
])
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]),
])
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,
)
# model =models.init_model(name=args.arch, num_classes = dataset.num_train_pids, loss = 'softmax')
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion_class = CrossEntropyLabelSmooth(num_classes=751)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
# optimizer = init_optim(args.optim, nn.Sequential([
# model.conv1,
# model.conv2,
# ]))
if args.stepsize > 0:
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']
# Parallel
if use_gpu:
model = nn.DataParallel(model).cuda()
# model.module.parameters() !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
if args.evaluate:
print('Evaluate only!')
test(model, queryloader, galleryloader, use_gpu)
return 0
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print('==>start training')
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
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(
) ### use_gpu .module. !!!!!!!!
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')) # fpath=/log/checkpoint_ep().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 main():
use_gpu = torch.cuda.is_available()
# use_gpu = False
if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else False
if not args.test:
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_img_dataset(
root=args.root, name=args.dataset, split_id=args.split_id,
)
print('dataset',dataset)
# data augmentation
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
# T.Resize(size=(384,128),interpolation=3),
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.Resize(size=(384,128),interpolation=3),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
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,
)
print('len of trainloader',len(trainloader))
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test,train=False),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
print('len of queryloader',len(queryloader))
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test,train=False),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
#embed()
print('len of galleryloader',len(galleryloader))
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_vids,
loss={'softmax','metric'}, aligned =True, use_gpu=use_gpu)
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0))
print('Model ',model)
print('num_classes',dataset.num_train_vids)
if args.labelsmooth:
criterion_class = CrossEntropyLabelSmooth(num_classes=dataset.num_train_vids, use_gpu=use_gpu)
else:
# criterion_class = CrossEntropyLoss(use_gpu=use_gpu)
criterion_class = nn.CrossEntropyLoss()
criterion_metric = TripletLossAlignedReID(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
if args.stepsize > 0:
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.test:
print("test aicity dataset")
if args.use_track_info:
g_track_id = get_track_id(args.root)
test(model, queryloader, galleryloader, use_gpu,dataset_q=dataset.query,dataset_g=dataset.gallery,track_id_tmp=g_track_id,rank=100)
else:
test(model, queryloader, galleryloader, use_gpu,dataset_q=dataset.query,dataset_g=dataset.gallery,rank=100)
return 0
def main():
# 判断是否使用gpu
use_gpu = torch.cuda.is_available()
# 若指定只使用cpu，则置use_gpu = False
if args.use_cpu: use_gpu = False
pin_memory = True if use_gpu else False # 避免内存浪费
# 日志打印设置
# 训练阶段存放在log_train.txt，测试阶段存放在log_test.txt
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))
# 若使用gpu，进行相关优化设置
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
dataset = data_manager.init_img_dataset(root=args.root,
name=args.dataset,
split_id=args.split_id)
# 创建dataloader & 进行 augmentation，训练时进行数据增广，测试时不需要
# 3个data_loader train query gallery
# 训练用transform
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]), #归一化
])
# 测试用transform
transform_test = T.Compose([
T.Resize((args.height, args.width)), #只做resize处理，将图像统一到同一尺寸
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 导入dataset_loader
# param@drop_last:把尾部多余数据扔掉，不用做训练了
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,
)
# param@drop_last：在test阶段，每一个样本都不能丢
# param@shuffle：在test阶段就不要打乱了
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,
)
# 初始化模型Resnet50
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'softmax', 'metric'},
aligned=True,
use_gpu=use_gpu)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
# 确定损失类型
criterion_class = CrossEntropyLoss(use_gpu=use_gpu)
# 优化器
# [email protected]()：更新模型的所有参数 若更新某一层：model.conv1 model.fc；若更新两层：nn.Sequential([model.conv1,model.conv2])
# param@lr：learning rate，学习率
# param@decay：模型的正则化参数
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
# 学习率衰减，逐步减小步长，采用阶梯型衰减
# param@gamma：衰减倍率
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
# 设置开始训练的epoch
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 0
print('start training')
# 开始进行训练
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, optimizer, trainloader, use_gpu)
# save_checkpoint是个字典dic
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
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, hflip_prob=0.5):
trans = [T.ToTensor()]
if hflip_prob > 0:
trans.append(T.RandomHorizontalFlip(hflip_prob))
self.transforms = T.Compose(trans)
def main():
# 是否使用GPU和是否节省显存
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
pin_memory = True if use_gpu else False
# Log文件的输出
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))
# GPU调用
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
cudnn.benchmark = True # 表示使用cudnn
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU !")
# 初始化dataset
dataset = data_manager.Market1501(root=args.root)
# dataloader(train query gallery) 和 增强
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(p=0.5),
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])
])
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
sampler=RandomIdentitySampler(dataset.train,
num_instance=args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=
True # 训练时会存在epoch % batchsize != 0 的情况，那么余数的图片是否还要训练？drop_last= True就是舍去这些图片
)
queryloader = DataLoader(ImageDataset(dataset.query,
transform=transform_test),
batch_size=args.train_batch,
num_workers=args.workers,
shuffle=False,
pin_memory=pin_memory)
galleryloader = DataLoader(ImageDataset(dataset.gallery,
transform=transform_test),
batch_size=args.train_batch,
num_workers=args.workers,
shuffle=False,
pin_memory=pin_memory)
# 加载模型
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'softmax', 'metric'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
# 损失和优化器
criterion_class = nn.CrossEntropyLoss()
criterion_metric = TripletLoss(margin=args.margin)
# optimizer = torch.optim.adam()
# 只更新其中某两层(先运行下行语句看看要更新哪几层)
# print(*list(model.children()))
# optimizer = init_optim(args.optim, model.parameters(nn.Sequential([
# *list(model.children())[:-2]
# ])))
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
if args.stepsize > 0:
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 0
start_time = time.time()
train_time = 0
# 训练
print("Start Traing!")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, criterion_metric, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
# 测试以及存模型
# step()了才会衰减
if args.stepsize > 0:
scheduler.step()
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))
