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'))
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_vidreid_dataset(root=args.root,
name=args.dataset)
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]),
])
pin_memory = True if use_gpu else False
# decompose tracklets into images
new_train = []
for img_paths, pid, camid in dataset.train:
for img_path in img_paths:
new_train.append((img_path, pid, camid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
VideoDataset(dataset.query,
seq_len=args.seq_len,
sample='evenly',
transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
VideoDataset(dataset.gallery,
seq_len=args.seq_len,
sample='evenly',
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'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids,
use_gpu=use_gpu)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
if args.fixbase_epoch > 0:
optimizer_tmp = init_optim(args.optim, model.classifier.parameters(),
args.fixbase_lr, args.weight_decay)
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, args.pool, use_gpu)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
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(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 (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, args.pool, 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():
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'))
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.Resize((args.height, args.width)),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop([args.height, args.width]),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
torchvision.transforms.RandomErasing(p=0.5,
scale=(0.02, 0.4),
ratio=(0.3, 3.33),
value=(0.4914, 0.4822, 0.4465))
# T.RandomErasing(probability=0.5, sh=0.4, mean=(0.4914, 0.4822, 0.4465)),
])
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
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={'cent'})
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_cent = CenterLoss(num_classes=dataset.num_train_pids,
feat_dim=model.feat_dim,
use_gpu=use_gpu)
optimizer_model = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
optimizer_cent = torch.optim.SGD(criterion_cent.parameters(),
lr=args.lr_cent)
#only the optimizer_model use learning rate schedule
# if args.stepsize > 0:
# scheduler = lr_scheduler.StepLR(optimizer_model, step_size=args.stepsize, gamma=args.gamma)
'''------Modify lr_schedule here------'''
current_schedule = init_lr_schedule(schedule=args.schedule,
warm_up_epoch=args.warm_up_epoch,
half_cos_period=args.half_cos_period,
lr_milestone=args.lr_milestone,
gamma=args.gamma,
stepsize=args.stepsize)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer_model,
lr_lambda=current_schedule)
'''------Please refer to the args.xxx for details of hyperparams------'''
# embed()
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
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_xent, criterion_cent, optimizer_model,
optimizer_cent, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.schedule: 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 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))
return coloured_mask
img_paths = glob.glob("../main20170707/org_imgs/*.png")
# img_paths = glob.glob("../data/tool/org_imgs/*.png")
# dataset = Dataset(img_paths, annotation, is_train=False)
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()
confidence = 0.5
# for idx in tqdm.tqdm(range(dataset.__len__())):
for idx in tqdm(range(len(img_paths))):
# Prediction
# img, _ = dataset.__getitem__(idx)
img_path = img_paths[idx]
img = Image.open(img_path)
transform = T.Compose([T.ToTensor()])
img = torchvision.transforms.functional.to_tensor(img)
pred = model([img.to(device)])
pred_score = list(pred[0]['scores'].detach().cpu().numpy())
pred_t = [
pred_score.index(x) for x in pred_score if x > confidence
]
masks = (pred[0]['masks'] > 0.5).squeeze().detach().cpu().numpy()
if masks.ndim == 2:
masks = masks.reshape([1, masks.shape[0], masks.shape[1]])
pred_class = [
CLASS_NAMES[i] for i in list(pred[0]['labels'].cpu().numpy())
]
pred_boxes = [[
(i[0], i[1]), (i[2], i[3])
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
from torch.autograd import Variable
import transforms as transforms
from skimage import io
from skimage.transform import resize
from models import *
import cv2 as cv
import dlib
import time
transform_test = transforms.Compose([
transforms.ToTensor(),
])
def rect_to_bb(rect):
x = rect.left()
y = rect.top()
w = rect.right() - x
h = rect.bottom() - y
# return a tuple of (x, y, w, h)
return x, y, w, h
def apply_offsets(face_coordinates, offsets):
x, y, width, height = face_coordinates
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'))
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_dataset(root=args.root,
name=args.dataset,
cls_sample=args.cls_sample)
# print(dataset.train)
# print(1)
# 解释器：创建一个transform处理图像数据的设置
# T.Random2DTranslation：随机裁剪
# T.RandomHorizontalFlip: 给定概率进行随机水平翻转
# T.ToTensor: 将PIL或numpy向量[0,255]=>tensor[0.0,1.0]
# T.Normalize：用均值和标准偏差标准化张量图像，mean[ , , ]三个参数代表三通道
transform_train = T.Compose([
# T.RandomCrop(224),
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(256),
# T.CenterCrop(224),
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
m = dataset.train
print(1)
# Dataloader 提供队列和线程
# ImageDataset:return data =>img, pid, camid
# RandomIdentitySampler:定义从数据集中抽取样本的策略
# num_workers: 子进程数
# print(dataset.train)
trainloader = DataLoader(
AGE_Gender_ImageDataset(dataset.train, transform=transform_train),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
testloader = DataLoader(
AGE_Gender_ImageDataset(dataset.test, 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.train_num_class,
loss={'xent'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
# criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.train_num_class, use_gpu=use_gpu)
age_criterion_xent = nn.CrossEntropyLoss()
gender_criterion_xent = nn.CrossEntropyLoss()
# optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=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, testloader, use_gpu)
return
start_time = time.time()
train_time = 0
# best_rank1 = -np.inf
best_score = 0
best_MAE = 0
best_gender_acc = 0
best_epoch = 0
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, age_criterion_xent, gender_criterion_xent,
optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Test")
MAE, Gender_acc = test(model, testloader, use_gpu)
Score = Gender_acc * 100 - MAE
is_best = Score > best_score
if is_best:
best_score = Score
best_MAE = MAE
best_gender_acc = Gender_acc
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': Score,
'epoch': epoch,
}, is_best,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print(
"==> Best best_score(Gender_acc-MAE) {} |Gender_acc {}\t MAE {}|achieved at epoch {}"
.format(best_score, best_gender_acc, best_MAE, 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))
"""Keep reading image until succeed.
This can avoid IOError incurred by heavy IO process."""
got_img = False
if not osp.exists(img_path):
raise IOError("{} does not exist".format(img_path))
while not got_img:
try:
img = Image.open(img_path).convert('RGB')
got_img = True
except IOError:
print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
pass
return img
transform_test = T.Compose([
T.Resize((args.height, args.width), interpolation=3),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def train(epoch,model_edge, model, model_tradclass,weight_softmax, criterion, optimizer, trainloader, learning_rate, use_gpu):
if not os.path.isdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_1"):
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_1")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_2")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_3")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_4")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_5")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_6")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_7")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_8")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_9")
os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_10")
def testseq(dataset_name, use_gpu):
dataset_root = './video2img/track1_sct_img_test_big/'
dataset = Graph_data_manager.AICityTrack2(root=dataset_root)
width = 224
height = 224
transform_train = T.Compose([
T.Random2DTranslation(height, 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((height, 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
seq_len = 4
num_instance = 4
train_batch = 32
test_batch = 1
queryloader = DataLoader(
VideoDataset(dataset.query,
seq_len=seq_len,
sample='dense',
transform=transform_test),
batch_size=test_batch,
shuffle=False,
num_workers=4,
pin_memory=pin_memory,
drop_last=False,
)
arch = "resnet50ta"
pretrained_model = "./log/track12_ta224_checkpoint_ep500.pth.tar"
start_epoch = 0
print("Initializing model: {}".format(arch))
dataset.num_train_pids = 517
if arch == 'resnet503d':
model = resnet3d.resnet50(num_classes=dataset.num_train_pids,
sample_width=width,
sample_height=height,
sample_duration=seq_len)
if not os.path.exists(pretrained_model):
raise IOError(
"Can't find pretrained model: {}".format(pretrained_model))
print("Loading checkpoint from '{}'".format(pretrained_model))
checkpoint = torch.load(pretrained_model)
state_dict = {}
for key in checkpoint['state_dict']:
if 'fc' in key: continue
state_dict[key.partition("module.")
[2]] = checkpoint['state_dict'][key]
model.load_state_dict(state_dict, strict=False)
else:
if not os.path.exists(pretrained_model):
model = models.init_model(name=arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
else:
model = models.init_model(name=arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
checkpoint = torch.load(pretrained_model)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch'] + 1
print("Loaded checkpoint from '{}'".format(pretrained_model))
print("- start_epoch: {}\n- rank1: {}".format(
start_epoch, checkpoint['rank1']))
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=0.3)
lr = 0.0003
gamma = 0.1
stepsize = 200
weight_decay = 5e-04
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
if stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=stepsize,
gamma=gamma)
start_epoch = start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
test(model, queryloader, 'avg', use_gpu, dataset, -1, meta_data_tab=None)
def main(parser_data):
device = torch.device(
parser_data.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
data_transform = {"val": transforms.Compose([transforms.ToTensor()])}
# read class_indict
label_json_path = './pascal_voc_classes.json'
assert os.path.exists(
label_json_path), "json file {} dose not exist.".format(
label_json_path)
json_file = open(label_json_path, 'r')
class_dict = json.load(json_file)
category_index = {v: k for k, v in class_dict.items()}
VOC_root = parser_data.data_path
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError(
"VOCdevkit dose not in path:'{}'.".format(VOC_root))
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = parser_data.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
# load validation data set
val_data_set = VOC2007DataSet(VOC_root, data_transform["val"], "val.txt")
val_data_set_loader = torch.utils.data.DataLoader(
val_data_set,
batch_size=batch_size,
shuffle=False,
num_workers=nw,
collate_fn=val_data_set.collate_fn)
# create model num_classes equal background + 20 classes
backbone = resnet50_fpn_backbone()
model = FasterRCNN(backbone=backbone,
num_classes=parser_data.num_classes + 1)
# 载入你自己训练好的模型权重
weights_path = parser_data.weights
assert os.path.exists(weights_path), "not found {} file.".format(
weights_path)
weights_dict = torch.load(weights_path, map_location=device)
model.load_state_dict(weights_dict['model'])
# print(model)
model.to(device)
# evaluate on the test dataset
coco = get_coco_api_from_dataset(val_data_set)
iou_types = ["bbox"]
coco_evaluator = CocoEvaluator(coco, iou_types)
cpu_device = torch.device("cpu")
model.eval()
with torch.no_grad():
for image, targets in tqdm(val_data_set_loader, desc="validation..."):
# 将图片传入指定设备device
image = list(img.to(device) for img in image)
# inference
outputs = model(image)
outputs = [{k: v.to(cpu_device)
for k, v in t.items()} for t in outputs]
res = {
target["image_id"].item(): output
for target, output in zip(targets, outputs)
}
coco_evaluator.update(res)
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
coco_eval = coco_evaluator.coco_eval["bbox"]
# calculate COCO info for all classes
coco_stats, print_coco = summarize(coco_eval)
# calculate voc info for every classes(IoU=0.5)
voc_map_info_list = []
for i in range(len(category_index)):
stats, _ = summarize(coco_eval, catId=i)
voc_map_info_list.append(" {:15}: {}".format(category_index[i + 1],
stats[1]))
print_voc = "\n".join(voc_map_info_list)
print(print_voc)
# 将验证结果保存至txt文件中
with open("record_mAP.txt", "w") as f:
record_lines = [
"COCO results:", print_coco, "", "mAP(IoU=0.5) for each category:",
print_voc
]
f.write("\n".join(record_lines))
def main(**kwargs):
training_path = kwargs.get('training_data_path')
checkpoint_path = kwargs.get('checkpoint_path')
tensorboard_log_dir = kwargs.get('tensorboard_log_dir')
if not os.path.isdir(checkpoint_path):
os.mkdir(checkpoint_path)
backbone_name = kwargs.get('backbone')
criterion_name = kwargs.get('criterion').upper()
optimizer_name = kwargs.get('optimizer').upper()
scheduler = kwargs.get('scheduler',None)
pretrained = kwargs.get('pretrained')
num_classes = kwargs.get('num_classes')
device = kwargs.get('device')
batch_size = kwargs.get('batch_size')
epochs = kwargs.get('epochs')
hyperparameters = kwargs.get('hyperparameters',{})
augmentations = kwargs.get('augmentations',{})
verbose = kwargs.get('verbose')
train_split = kwargs.get('train_split')
nfolds = kwargs.get('nfolds')
val_splits = [(1-train_split) / nfolds] * nfolds
resume = kwargs.get('resume')
only_weights = kwargs.get('only_weights')
seed = hyperparameters.get('seed')
random_jitter = augmentations.get('jitter',{})
random_horizontal_flip = augmentations.get('horizontal_flip', 0.5)
random_rotation = augmentations.get('rotation', 20)
writer = SummaryWriter(log_dir=tensorboard_log_dir, flush_secs=20)
if seed: seed_everything(seed)
# TODO calculate mean and std
mean = hyperparameters.get('mean',0)
std = hyperparameters.get('std',1)
splits = [train_split]+val_splits
assert sum(splits) <= 1,"given splits must be lower or equal than 1"
original_img_size = 96
criterion = get_criterion(criterion_name)
backbone = get_backbone(backbone_name, pretrained=pretrained)
model = FacialKeypointsDetector(backbone, criterion=criterion,
device=device, num_classes=num_classes)
optimizer = get_optimizer(optimizer_name, model.parameters(),
kwargs=hyperparameters.get('optimizer',{}))
scaler = GradScaler()
val_transforms = None
val_target_transform = TargetTransform(original_img_size)
train_transform = train_target_transform = None
train_transforms = transforms.TrainTransforms(model.get_input_size(), original_img_size,
mean=mean, std=std, brightness=random_jitter.get('brightness'),
contrast=random_jitter.get('contrast'), saturation=random_jitter.get('saturation'),
hue=random_jitter.get('hue'), rotation_degree=random_rotation,
hflip=random_horizontal_flip)
val_transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(model.get_input_size()),
transforms.ToTensor(),
transforms.Normalize(mean,std)])
train_dataset,*val_datasets = get_training_datasets(root_path=training_path,
train_transforms=(train_transforms,train_transform,train_target_transform),
val_transforms=(val_transforms,val_transform,val_target_transform),
split_ratios=splits)
val_dls = []
train_dl = torch.utils.data.DataLoader(train_dataset,
num_workers=4, batch_size=batch_size,
pin_memory=True, collate_fn=custom_collate_fn, shuffle=True)
for val_ds in val_datasets:
val_dls.append( torch.utils.data.DataLoader(
val_ds, batch_size=batch_size, num_workers=2) )
current_epoch = 0
best_loss = math.inf
if resume:
print(Fore.CYAN, f"loading checkpoint from {checkpoint_path}",Style.RESET_ALL)
best_loss,current_epoch = load_checkpoint(model, optimizer, scheduler=scheduler,
save_path=checkpoint_path, suffix='last', only_weights=only_weights)
try:
for epoch in range(current_epoch,epochs):
training_loop(train_dl, model, epoch, epochs, optimizer, writer,scaler,
scheduler=scheduler, verbose=verbose)
val_losses = []
for i,val_dl in enumerate(val_dls):
val_loss = validation_loop(val_dl, model)
val_losses.append(val_loss)
print(Fore.LIGHTBLUE_EX, f"validation [{i+1}] loss: {val_loss:.07f}",Style.RESET_ALL)
writer.add_scalar(f'Loss/val_{i+1}', val_loss, epoch)
mean_val_loss = sum(val_losses) / len(val_losses)
print(Fore.LIGHTBLUE_EX, f"validation [mean] loss: {mean_val_loss:.07f}",Style.RESET_ALL)
writer.add_scalar(f'Loss/val_mean', mean_val_loss, epoch)
writer.flush()
if mean_val_loss < best_loss:
best_loss = mean_val_loss
print(Fore.CYAN, "saving best checkpoint...",Style.RESET_ALL)
save_checkpoint(model,optimizer,epoch,best_loss,
scheduler=scheduler, suffix='best', save_path=checkpoint_path)
print(Fore.CYAN, "saving last checkpoint...",Style.RESET_ALL)
save_checkpoint(model,optimizer,epoch,best_loss,
scheduler=scheduler, suffix='last', save_path=checkpoint_path)
except KeyboardInterrupt:
print(Fore.RED, "training interrupted with ctrl+c saving current state of the model",Style.RESET_ALL)
save_checkpoint(model,optimizer,epoch,best_loss,
scheduler=scheduler, suffix='last', save_path=checkpoint_path)
writer.close()
def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# Data
print('==> Preparing dataset %s' % args.dataset)
transform_train = transforms.Compose([
transforms.RandomCrop(28, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, )),
transforms.RandomErasing(probability=args.p,
sh=args.sh,
r1=args.r1,
mean=[0.4914]),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, )),
])
if args.dataset == 'fashionmnist':
dataloader = datasets.FashionMNIST
num_classes = 10
trainset = dataloader(root='../../dataset',
train=True,
download=True,
transform=transform_train)
trainloader = data.DataLoader(trainset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers)
testset = dataloader(root='../../dataset',
train=False,
download=False,
transform=transform_test)
testloader = data.DataLoader(testset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers)
# Model
print("==> creating model '{}'".format(args.arch))
if args.arch.startswith('wrn'):
model = models.__dict__[args.arch](
num_classes=num_classes,
depth=args.depth,
widen_factor=args.widen_factor,
dropRate=args.drop,
)
elif args.arch.endswith('resnet'):
model = models.__dict__[args.arch](
num_classes=num_classes,
depth=args.depth,
)
#model = torch.nn.DataParallel(model).cuda()
model = model.to(device)
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Resume
title = 'fashionmnist-' + args.arch
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
args.checkpoint = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
if args.evaluate:
print('\nEvaluation only')
#test_loss, test_acc = test(testloader, model, criterion, start_epoch, use_cuda)
#print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc))
model.eval()
id = 0
preds_list = []
with torch.no_grad():
for X, y in testloader:
batch_pred = list(
model(X.to(device)).argmax(dim=1).cpu().numpy())
for y_pred in batch_pred:
preds_list.append((id, y_pred))
id += 1
print('生成提交结果文件')
with open('submission.csv', 'w') as f:
f.write('ID,Prediction\n')
for id, pred in preds_list:
f.write('{},{}\n'.format(id, pred))
return
# Train and val
for epoch in range(start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['lr']))
train_loss, train_acc = train(trainloader, model, criterion, optimizer,
epoch, use_cuda)
test_loss, test_acc = test(testloader, model, criterion, epoch,
use_cuda)
# append logger file
logger.append(
[state['lr'], train_loss, test_loss, train_acc, test_acc])
# save model
is_best = test_acc > best_acc
best_acc = max(test_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
checkpoint=args.checkpoint)
logger.close()
logger.plot()
savefig(os.path.join(args.checkpoint, 'log.eps'))
print('Best acc:')
print(best_acc)
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 用来保存coco_info的文件
results_file = "results{}.txt".format(
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists("save_weights"):
os.makedirs("save_weights")
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
VOC_root = "./"
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError(
"VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"],
"train.txt")
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = 8
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_size=batch_size,
shuffle=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# load validation data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt")
val_data_set_loader = torch.utils.data.DataLoader(
val_data_set,
batch_size=batch_size,
shuffle=False,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# create model num_classes equal background + 20 classes
model = create_model(num_classes=21)
# print(model)
model.to(device)
train_loss = []
learning_rate = []
val_map = []
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# first frozen backbone and train 5 epochs #
# 首先冻结前置特征提取网络权重（backbone），训练rpn以及最终预测网络部分 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
for param in model.backbone.parameters():
param.requires_grad = False
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
init_epochs = 5
for epoch in range(init_epochs):
# train for one epoch, printing every 10 iterations
mean_loss, lr = utils.train_one_epoch(model,
optimizer,
train_data_loader,
device,
epoch,
print_freq=50)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [
str(round(i, 4)) for i in coco_info + [mean_loss.item(), lr]
]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
torch.save(model.state_dict(), "./save_weights/pretrain.pth")
# # # # # # # # # # # # # # # # # # # # # # # # # # # #
# second unfrozen backbone and train all network #
# 解冻前置特征提取网络权重（backbone），接着训练整个网络权重 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # #
# 冻结backbone部分底层权重
for name, parameter in model.backbone.named_parameters():
split_name = name.split(".")[0]
if split_name in ["0", "1", "2", "3"]:
parameter.requires_grad = False
else:
parameter.requires_grad = True
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.33)
num_epochs = 20
for epoch in range(init_epochs, num_epochs + init_epochs, 1):
# train for one epoch, printing every 50 iterations
mean_loss, lr = utils.train_one_epoch(model,
optimizer,
train_data_loader,
device,
epoch,
print_freq=50)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [
str(round(i, 4)) for i in coco_info + [mean_loss.item(), lr]
]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
# save weights
# 仅保存最后5个epoch的权重
if epoch in range(num_epochs + init_epochs)[-5:]:
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch
}
torch.save(save_files,
"./save_weights/mobile-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
def main():
torch.backends.cudnn.benchmark = True
args = parser.parse_args()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((224, 224)),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
### data loader => MJJ
checkpoint = torch.load(args.resnet_checkpoint)
state_dict = {
str.replace(k, 'module.', ''): v
for k, v in checkpoint['state_dict'].items()
}
fe = resnet.resnet50()
fe.load_state_dict(state_dict, False)
model = torch.nn.Sequential(fe, clsf())
checkpoint = (torch.load(args.clsf_checkpoint))
state_dict = {
str.replace(k, 'module.', ''): v
for k, v in checkpoint['model'].items()
}
model.load_state_dict(state_dict, False)
model = torch.nn.DataParallel(model).cuda()
model.eval()
cap = cv2.VideoCapture(args.input_filename)
fps = 1
n = 0
num_frames = 0
out_dict = []
check = 0
frame_counter = 0
images = []
with torch.no_grad():
while cap.isOpened():
ret, image = cap.read()
frame_counter += 1
if frame_counter >= (cap.get(cv2.CAP_PROP_FPS) / fps):
if not ret:
check = 10 #End
else:
image = Image.fromarray(image[:, :, ::-1])
images.append(image)
check += 1 #Append 10 times.
if check == 10:
check = 0
data = transform(images).unsqueeze(0).pin_memory().cuda(
non_blocking=True)
output = model(
data) #data : B*T*C*H*W // B : 1 // T : 10 frames.
y_friends = F.softmax(output, 1)
top5_value_friends, top5_index_friends = y_friends.topk(5)
top5_value_friends, top5_label_friends = top5_value_friends.tolist(
), [class_friends[i] for i in top5_index_friends[:, 0]]
for i in range(len(images)):
num_frames += 1
out_dict.append({
"type":
"location",
"class":
top5_label_friends[i],
"seconds":
float(num_frames) * 1.0 / float(fps)
})
images = []
if not ret:
break
frame_counter = 0
n += 1
if n % 100 == 0:
print('Processed {}/{} frames'.format(
n, int(cap.get(cv2.CAP_PROP_FRAME_COUNT))))
sys.stdout.flush()
if n == 1e10:
break
with jsonlines.open(args.output_filename + '.jsonl', mode='w') as writer:
writer.write_all(out_dict)
cap.release()
print('Done')
def load_data(test_data_path, preprocess_path, fold, batch_size, num_workers):
test_data_path = '/data/xxx/LUNA/rowfile/subset'
crop_size = 32
black_list = []
preprocess_path = '/data/xxx/LUNA/cls/crop_v3'
pix_value, npix = 0, 0
for file_name in os.listdir(preprocess_path):
if file_name.endswith('.npy'):
if file_name[:-4] in black_list:
continue
data = np.load(os.path.join(preprocess_path, file_name))
pix_value += np.sum(data)
npix += np.prod(data.shape)
pix_mean = pix_value / float(npix)
pix_value = 0
for file_name in os.listdir(preprocess_path):
if file_name.endswith('.npy'):
if file_name[:-4] in black_list: continue
data = np.load(os.path.join(preprocess_path, file_name)) - pix_mean
pix_value += np.sum(data * data)
pix_std = np.sqrt(pix_value / float(npix))
print(pix_mean, pix_std)
transform_train = transforms.Compose([
# transforms.RandomScale(range(28, 38)),
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.RandomYFlip(),
transforms.RandomZFlip(),
transforms.ZeroOut(4),
transforms.ToTensor(),
transforms.Normalize(
(pix_mean),
(pix_std)), # need to cal mean and std, revise norm func
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((pix_mean), (pix_std)),
])
# load data list
test_file_name_list = []
test_label_list = []
test_feat_list = []
data_frame = pd.read_csv('./data/annotationdetclsconvfnl_v3.csv',
names=[
'seriesuid', 'coordX', 'coordY', 'coordZ',
'diameter_mm', 'malignant'
])
all_list = data_frame['seriesuid'].tolist()[1:]
label_list = data_frame['malignant'].tolist()[1:]
crdx_list = data_frame['coordX'].tolist()[1:]
crdy_list = data_frame['coordY'].tolist()[1:]
crdz_list = data_frame['coordZ'].tolist()[1:]
dim_list = data_frame['diameter_mm'].tolist()[1:]
# test id
test_id_list = []
for file_name in os.listdir(test_data_path + str(fold) + '/'):
if file_name.endswith('.mhd'):
test_id_list.append(file_name[:-4])
mxx = mxy = mxz = mxd = 0
for srsid, label, x, y, z, d in zip(all_list, label_list, crdx_list,
crdy_list, crdz_list, dim_list):
mxx = max(abs(float(x)), mxx)
mxy = max(abs(float(y)), mxy)
mxz = max(abs(float(z)), mxz)
mxd = max(abs(float(d)), mxd)
if srsid in black_list:
continue
# crop raw pixel as feature
data = np.load(os.path.join(preprocess_path, srsid + '.npy'))
bgx = int(data.shape[0] / 2 - crop_size / 2)
bgy = int(data.shape[1] / 2 - crop_size / 2)
bgz = int(data.shape[2] / 2 - crop_size / 2)
data = np.array(data[bgx:bgx + crop_size, bgy:bgy + crop_size,
bgz:bgz + crop_size])
y, x, z = np.ogrid[-crop_size / 2:crop_size / 2,
-crop_size / 2:crop_size / 2,
-crop_size / 2:crop_size / 2]
mask = abs(y**3 + x**3 + z**3) <= abs(float(d))**3
feat = np.zeros((crop_size, crop_size, crop_size), dtype=float)
feat[mask] = 1
if srsid.split('-')[0] in test_id_list:
test_file_name_list.append(srsid + '.npy')
test_label_list.append(int(label))
test_feat_list.append(feat)
for idx in range(len(test_feat_list)):
test_feat_list[idx][-1] /= mxd
test_set = lunanod(preprocess_path,
test_file_name_list,
test_label_list,
test_feat_list,
train=False,
download=True,
transform=transform_test)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
return test_loader
def main(args):
init_distributed_mode(args)
print(args)
device = torch.device(args.device)
# 用来保存coco_info的文件
results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
# Data loading code
print("Loading data")
data_transform = {
"train": transforms.Compose([transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val": transforms.Compose([transforms.ToTensor()])
}
VOC_root = args.data_path
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_dataset = VOCDataSet(VOC_root, "2012", data_transform["train"], "train.txt")
# load validation data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_dataset = VOCDataSet(VOC_root, "2012", data_transform["val"], "val.txt")
print("Creating data loaders")
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
test_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
else:
train_sampler = torch.utils.data.RandomSampler(train_dataset)
test_sampler = torch.utils.data.SequentialSampler(val_dataset)
if args.aspect_ratio_group_factor >= 0:
# 统计所有图像比例在bins区间中的位置索引
group_ids = create_aspect_ratio_groups(train_dataset, k=args.aspect_ratio_group_factor)
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)
else:
train_batch_sampler = torch.utils.data.BatchSampler(
train_sampler, args.batch_size, drop_last=True)
data_loader = torch.utils.data.DataLoader(
train_dataset, batch_sampler=train_batch_sampler, num_workers=args.workers,
collate_fn=train_dataset.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
val_dataset, batch_size=1,
sampler=test_sampler, num_workers=args.workers,
collate_fn=train_dataset.collate_fn)
print("Creating model")
# create model
# 注意：不包含背景
model = create_model(num_classes=args.num_classes)
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
model_without_ddp = model.module
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
# 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练
if args.resume:
# If map_location is missing, torch.load will first load the module to CPU
# and then copy each parameter to where it was saved,
# which would result in all processes on the same machine using the same set of devices.
checkpoint = torch.load(args.resume, map_location='cpu') # 读取之前保存的权重文件(包括优化器以及学习率策略)
model_without_ddp.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
if args.test_only:
utils.evaluate(model, data_loader_test, device=device)
return
train_loss = []
learning_rate = []
val_map = []
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
mean_loss, lr = utils.train_one_epoch(model, optimizer, data_loader, device,
epoch, args.print_freq, warmup=True)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update learning rate
lr_scheduler.step()
# evaluate after every epoch
coco_info = utils.evaluate(model, data_loader_test, device=device)
val_map.append(coco_info[1]) # pascal mAP
# 只在主进程上进行写操作
if args.rank in [-1, 0]:
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [str(round(i, 4)) for i in coco_info + [mean_loss.item()]] + [str(round(lr, 6))]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
if args.output_dir:
# 只在主节点上执行保存权重操作
save_on_master({
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'args': args,
'epoch': epoch},
os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
if args.rank in [-1, 0]:
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
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'))
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_dataset(name=args.dataset)
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]),
])
pin_memory = False
trainloader = DataLoader(
VideoDataset(dataset.train,
seq_len=args.seq_len,
sample='random',
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(
VideoDataset(dataset.query,
seq_len=args.seq_len,
sample='dense',
transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
VideoDataset(dataset.gallery,
seq_len=args.seq_len,
sample='dense',
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', 'htri'})
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 = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
start_time = time.time()
print(start_time)
for batch_idx, (imgs, pids, _) in enumerate(trainloader):
print(batch_idx)
print('x')
if use_gpu:
imgs, pids = imgs.cuda(), pids.cuda()
imgs, pids = Variable(imgs), Variable(pids)
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
def main(args):
print(args)
# mp.spawn(main_worker, args=(args,), nprocs=args.world_size, join=True)
utils.init_distributed_mode(args)
device = torch.device(args.device)
# Data loading code
print("Loading data")
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
VOC_root = args.data_path
# load train data set
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], True)
# load validation data set
val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], False)
print("Creating data loaders")
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_data_set)
test_sampler = torch.utils.data.distributed.DistributedSampler(
val_data_set)
else:
train_sampler = torch.utils.data.RandomSampler(train_data_set)
test_sampler = torch.utils.data.SequentialSampler(val_data_set)
if args.aspect_ratio_group_factor >= 0:
# 统计所有图像比例在bins区间中的位置索引
group_ids = create_aspect_ratio_groups(
train_data_set, k=args.aspect_ratio_group_factor)
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids,
args.batch_size)
else:
train_batch_sampler = torch.utils.data.BatchSampler(train_sampler,
args.batch_size,
drop_last=True)
data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_sampler=train_batch_sampler,
num_workers=args.workers,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(val_data_set,
batch_size=1,
sampler=test_sampler,
num_workers=args.workers,
collate_fn=utils.collate_fn)
print("Creating model")
model = create_model(num_classes=21)
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
# 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练
if args.resume:
# If map_location is missing, torch.load will first load the module to CPU
# and then copy each parameter to where it was saved,
# which would result in all processes on the same machine using the same set of devices.
checkpoint = torch.load(
args.resume, map_location='cpu') # 读取之前保存的权重文件(包括优化器以及学习率策略)
model_without_ddp.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
if args.test_only:
utils.evaluate(model, data_loader_test, device=device)
return
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
utils.train_one_epoch(model, optimizer, data_loader, device, epoch,
args.print_freq)
lr_scheduler.step()
if args.output_dir:
# 只在主节点上执行保存权重操作
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'args': args,
'epoch': epoch
}, os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
# evaluate after every epoch
utils.evaluate(model, data_loader_test, device=device)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.cuda("cpu")
sys.stdout = Logger(osp.join(PATH, 'log_train.txt'))
print("Dataset is being initialized")
dataset = dataset_manager.init_img_dataset(root='data',
name=dataset_name,
split_id=0)
tfms_train = tfms.Compose([
tfms.Random2DTranslation(256, 128),
tfms.RandomHorizontalFlip(),
tfms.ToTensor(),
tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
tfms_test = tfms.Compose([
tfms.Resize((256, 128)),
tfms.ToTensor(),
tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True
trainloader = DataLoader(
ImageDataset(dataset.train, transform=tfms_train),
sampler=RandomIdentitySampler(dataset.train, num_instances=4),
batch_size=32,
def get_transform(is_train):
transforms = []
transforms.append(T.ToTensor())
return T.Compose(transforms)
def __init__(self):
self.transforms = T.ToTensor()
import torch.nn as nn
import torch.nn.functional as F
import os
from torch.autograd import Variable
import transforms as transforms
from skimage import io
from skimage.transform import resize
from models import *
cut_size = 44
transform_test = transforms.Compose([
transforms.TenCrop(cut_size),
transforms.Lambda(lambda crops: torch.stack(
[transforms.ToTensor()(crop) for crop in crops])),
])
def rgb2gray(rgb):
return np.dot(rgb[..., :3], [0.299, 0.587, 0.114])
raw_img = io.imread('images/1.jpg')
gray = rgb2gray(raw_img)
gray = resize(gray, (48, 48), mode='symmetric').astype(np.uint8)
img = gray[:, :, np.newaxis]
img = np.concatenate((img, img, img), axis=2)
img = Image.fromarray(img)
learning_rate_decay_rate = 0.9 # 0.9
cut_size = 44
total_epoch = 250
path = os.path.join(opt.dataset + '_' + opt.model)
dir_path = '/home/zjian30/test/Facial-Expression-Recognition.Pytorch/'
path = dir_path + path
print(path)
# Data
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.RandomCrop(44),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.TenCrop(cut_size),
transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),
])
trainset = FER2013(split = 'Training', transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.bs, shuffle=True, num_workers=1)
PublicTestset = FER2013(split = 'PublicTest', transform=transform_test)
PublicTestloader = torch.utils.data.DataLoader(PublicTestset, batch_size=opt.bs, shuffle=False, num_workers=1)
PrivateTestset = FER2013(split = 'PrivateTest', transform=transform_test)
PrivateTestloader = torch.utils.data.DataLoader(PrivateTestset, batch_size=opt.bs, shuffle=False, num_workers=1)
# Model
def attr_main():
stale_step = args.stalestep
runId = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
args.save_dir = os.path.join(args.save_dir, runId)
if not os.path.exists(args.save_dir):
os.mkdir(args.save_dir)
print(args.save_dir)
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('./log_train_' + runId + '.txt')
else:
sys.stdout = Logger('./log_test_' + runId + '.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_dataset(name=args.dataset, min_seq_len=args.seq_len, attr=True)
args.attr_lens = [[], [9, 10, 2]] #dataset.attr_lens
args.columns = ["upcolor", "downcolor", "gender"] #dataset.columns
print("Initializing model: {}".format(args.arch))
# if args.arch == "resnet50ta_attr" or args.arch == "resnet50ta_attr_newarch":
if args.arch == 'attr_resnet503d':
model = models.init_model(name=args.arch,
attr_lens=args.attr_lens,
model_type=args.model_type,
num_classes=dataset.num_train_pids,
sample_width=args.width,
sample_height=args.height,
sample_duration=args.seq_len)
torch.backends.cudnn.benchmark = False
else:
model = models.init_model(name=args.arch,
attr_lens=args.attr_lens,
model_type=args.model_type)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
if args.dataset == "duke":
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])
])
elif args.dataset == "mars":
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])
])
pin_memory = True if use_gpu else False
# if args.predict:
# attr_predict(model, transform_test, use_gpu)
# return
# # For SQL mmir, just first_img_path/ random_img_path sampling works
# # For other mmir, call first_img/ random_img sampling
# mmirImgQueryLoader = VideoDataset(dataset.mmir_img_query, seq_len=1, # dataset.mmir_img_query[0:2]
# sample='random_img_path', # sample='random', #changing to random to get the imgs for running attr. extraction model, does the img_path match with 'random_img_path' sampling -- yap, the samples returned with both sampling matches in postgres and here, so next just get those images and run the attribute ex. models on them
# transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# # The above function basically just tells how we sample dfrom the img list of this specific tracklet (has pid, camid, attrs)
# # sampling more from gallery as test and gallery has same person, train has different persons
# # during retrieval same person with same features should have higher rank than different person with same features
# # 2/3 for image, rest is for video
# mmirImgQueryLoader = VideoDataset(
# dataset.mmir_img_gallery,
# seq_len=1, # dataset.mmir_img_query[0:2]
# sample='random_img_path', # sample='random',
# transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# print(len(dataset.mmir_img_gallery))
# print(len(dataset.mmir_video_gallery))
# # import psycopg2
# # # set-up a postgres connection
# # conn = psycopg2.connect(database='ng', user='******',password='******',
# # host='146.148.89.5', port=5432)
# # dbcur = conn.cursor()
# # print("connection successful")
# # for batch_idx, (pids, camids, attrs, img_path) in enumerate(tqdm(mmirImgQueryLoader)):
# # #print(pids)
# # images = list()
# # images.append(img_path)
# # # just save the img in img_path
# # sql = dbcur.mogrify("""
# # INSERT INTO mmir_ground (
# # mgid,
# # ctype,
# # pid,
# # camid,
# # ubc,
# # lbc,
# # gender,
# # imagepaths
# # ) VALUES (DEFAULT, %s,%s,%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING;""", ("image",
# # str(pids),
# # str(camids),
# # str(attrs[0]),
# # str(attrs[1]),
# # str(attrs[2]), images)
# # )
# # # print(sql)
# # dbcur.execute(sql)
# # conn.commit()
# # print(dataset.mmir_video_query[0]) # if this works, for sdml or other mmir training will just take some part of train for img, some for video
# # For SQL mmir, just first_img_path sampling works, as not even using the image_array, using just the attributes
# # For other mmir, call random sampling
# # sampling the whole video has no impact in mmir, as we cannot take avg. of the separated lists attributes "accuracy"
# # we need frames of whose attributes would be one single one, and then we search in img and text that attributes
# mmirVideoQueryLoader = VideoDataset(dataset.mmir_video_query, seq_len=10, # 100 * 10 = 1000 frames as image
# sample='random_video', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# mmirVideoQueryLoader = VideoDataset(dataset.mmir_video_gallery, seq_len=6, # 100 * 10 = 1000 frames as image
# sample='random_video', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# # for batch_idx, (pids, camids, attrs, img_paths) in enumerate(tqdm(mmirVideoQueryLoader)):
# # #print(pids)
# # sql = dbcur.mogrify("""
# # INSERT INTO mmir_ground (
# # mgid,
# # ctype,
# # pid,
# # camid,
# # ubc,
# # lbc,
# # gender,
# # imagepaths
# # ) VALUES (DEFAULT, %s,%s,%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING;""", ("video",
# # str(pids),
# # str(camids),
# # str(attrs[0]),
# # str(attrs[1]),
# # str(attrs[2]), img_paths)
# # )
# # dbcur.execute(sql)
# # conn.commit()
# # conn.close()
# # if i want to add more mmir image, video samples,
# # 1. in line 116, pass num_mmir_query_imgs=1000, num_mmir_query_videos=100
# # 2. just uncomment the SQL query code, the first 1000 & 100 would be same, as seed is set, the later ones will be added in postgres
# # hahahah, way surpassed above steps, created new samples from query_id, appended them; then created samples from gallery_id, added them by calling sampler again
# trainloader = DataLoader(
# VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', transform=transform_train, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens),
# 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,
# )
# validloader = VideoDataset(dataset.valid, seq_len=args.seq_len,
# sample='dense', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# #queryloader = VideoDataset(dataset.query + dataset.gallery, seq_len=args.seq_len, sample='single' transform=transform_test, attr=True,
# queryloader = VideoDataset(dataset.query + dataset.gallery, seq_len=args.seq_len, #args.seq_len, sample='dense'
# sample='dense', transform=transform_test, attr=True,
# attr_loss=args.attr_loss, attr_lens=args.attr_lens)
# # queryLoaderMIT = VideoDataset(dataset.query + dataset.gallery, seq_len=1,
# # sample='random_img_path', transform=transform_test, attr=True,
# # attr_loss=args.attr_loss, attr_lens=args.attr_lens)
start_epoch = args.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
start_time = time.time()
if args.arch == 'resnet503d':
torch.backends.cudnn.benchmark = False
# print("Run attribute pre-training")
if args.attr_loss == "cropy":
criterion = nn.CrossEntropyLoss()
elif args.attr_loss == "mse":
criterion = nn.MSELoss()
if args.colorsampling:
attr_test_MIT(model, criterion, queryLoaderMIT, use_gpu)
return
if args.evaluate:
print("Evaluate only")
# model_root = "/data/chenzy/models/mars/2019-02-26_21-02-13"
model_root = args.eval_model_root
model_paths = []
for m in os.listdir(model_root):
if m.endswith("pth"):
model_paths.append(m)
model_paths = sorted(model_paths,
key=lambda a: float(a.split("_")[1]),
reverse=True)
#print(model_paths)
#input()
# model_paths = ['rank1_2.8755379380596713_checkpoint_ep500.pth']
for m in model_paths[0:1]:
model_path = os.path.join(model_root, m)
print(model_path)
old_weights = torch.load(model_path)
new_weights = model.module.state_dict()
for k in new_weights:
if k in old_weights:
new_weights[k] = old_weights[k]
model.module.load_state_dict(new_weights)
if args.predict:
attr_predict(model, transform_test, use_gpu)
return
avr_acc = attr_test(model, criterion, queryloader, use_gpu)
# break
# test(model, queryloader, galleryloader, args.pool, use_gpu)
return
if use_gpu:
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.module.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
# avr_acc = attr_test(model, criterion, queryloader, use_gpu)
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
best_avr = 0
no_rise = 0
for epoch in range(start_epoch, args.max_epoch):
print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch))
attr_train(model, criterion, optimizer, trainloader, use_gpu)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and ((epoch + 1) % (args.eval_step) == 0 or
(epoch + 1) == args.max_epoch):
# avr_acc = attr_test(model, criterion, queryloader, use_gpu)
avr_acc = attr_test(model, criterion, validloader, use_gpu)
print("avr", avr_acc)
if avr_acc > best_avr:
no_rise = 0
print("==> Test")
best_avr = avr_acc
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(
state_dict,
osp.join(
args.save_dir, "avr_" + str(avr_acc) +
'_checkpoint_ep' + str(epoch + 1) + '.pth'))
else:
no_rise += 1
print("no_rise:", no_rise)
if no_rise > stale_step:
break
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
self.transforms = T.Compose([
T.ToTensor(),
T.Normalize(mean=mean, std=std),
])
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists("save_weights"):
os.makedirs("save_weights")
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
VOC_root = "../../data_set/"
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError(
"VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], True)
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = 4
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_size=batch_size,
shuffle=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# load validation data set
val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], False)
val_data_set_loader = torch.utils.data.DataLoader(
val_data_set,
batch_size=batch_size,
shuffle=False,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# create model num_classes equal background + 20 classes
model = create_model(num_classes=21)
# print(model)
model.to(device)
train_loss = []
learning_rate = []
val_mAP = []
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# first frozen backbone and train 5 epochs #
# 首先冻结前置特征提取网络权重（backbone），训练rpn以及最终预测网络部分 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
for param in model.backbone.parameters():
param.requires_grad = False
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
num_epochs = 5
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
utils.train_one_epoch(model,
optimizer,
train_data_loader,
device,
epoch,
print_freq=50,
train_loss=train_loss,
train_lr=learning_rate)
# evaluate on the test dataset
utils.evaluate(model,
val_data_set_loader,
device=device,
mAP_list=val_mAP)
torch.save(model.state_dict(), "./save_weights/pretrain.pth")
# # # # # # # # # # # # # # # # # # # # # # # # # # # #
# second unfrozen backbone and train all network #
# 解冻前置特征提取网络权重（backbone），接着训练整个网络权重 #
# # # # # # # # # # # # # # # # # # # # # # # # # # # #
# 冻结backbone部分底层权重
for name, parameter in model.backbone.named_parameters():
split_name = name.split(".")[0]
if split_name in ["0", "1", "2", "3"]:
parameter.requires_grad = False
else:
parameter.requires_grad = True
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5,
gamma=0.33)
num_epochs = 20
for epoch in range(num_epochs):
# train for one epoch, printing every 50 iterations
utils.train_one_epoch(model,
optimizer,
train_data_loader,
device,
epoch,
print_freq=50,
train_loss=train_loss,
train_lr=learning_rate)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
utils.evaluate(model,
val_data_set_loader,
device=device,
mAP_list=val_mAP)
# save weights
if epoch > 10:
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch
}
torch.save(save_files,
"./save_weights/mobile-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_mAP) != 0:
from plot_curve import plot_map
plot_map(val_mAP)
def main(parser_data):
device = torch.device(
parser_data.device if torch.cuda.is_available() else "cpu")
print("Using {} device training.".format(device.type))
# 用来保存coco_info的文件
results_file = "results{}.txt".format(
datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
data_transform = {
"train":
transforms.Compose(
[transforms.ToTensor(),
transforms.RandomHorizontalFlip(0.5)]),
"val":
transforms.Compose([transforms.ToTensor()])
}
VOC_root = parser_data.data_path
# check voc root
if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
raise FileNotFoundError(
"VOCdevkit dose not in path:'{}'.".format(VOC_root))
# load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"],
"train.txt")
# 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
batch_size = parser_data.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(
train_data_set,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# load validation data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt")
val_data_set_loader = torch.utils.data.DataLoader(
val_data_set,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=train_data_set.collate_fn)
# create model num_classes equal background + 20 classes
model = create_model(num_classes=parser_data.num_classes + 1,
device=device)
# print(model)
model.to(device)
# define optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.33)
# 如果指定了上次训练保存的权重文件地址，则接着上次结果接着训练
if parser_data.resume != "":
checkpoint = torch.load(parser_data.resume, map_location=device)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
parser_data.start_epoch = checkpoint['epoch'] + 1
print("the training process from epoch{}...".format(
parser_data.start_epoch))
train_loss = []
learning_rate = []
val_map = []
for epoch in range(parser_data.start_epoch, parser_data.epochs):
# train for one epoch, printing every 10 iterations
mean_loss, lr = utils.train_one_epoch(model,
optimizer,
train_data_loader,
device=device,
epoch=epoch,
print_freq=50,
warmup=True)
train_loss.append(mean_loss.item())
learning_rate.append(lr)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
coco_info = utils.evaluate(model, val_data_set_loader, device=device)
# write into txt
with open(results_file, "a") as f:
# 写入的数据包括coco指标还有loss和learning rate
result_info = [
str(round(i, 4)) for i in coco_info + [mean_loss.item()]
] + [str(round(lr, 6))]
txt = "epoch:{} {}".format(epoch, ' '.join(result_info))
f.write(txt + "\n")
val_map.append(coco_info[1]) # pascal mAP
# save weights
save_files = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch
}
torch.save(save_files,
"./save_weights/resNetFpn-model-{}.pth".format(epoch))
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
from plot_curve import plot_loss_and_lr
plot_loss_and_lr(train_loss, learning_rate)
# plot mAP curve
if len(val_map) != 0:
from plot_curve import plot_map
plot_map(val_map)
def from_image(cls, img, sr_factor):
pil = transforms.ToPILImage()(transforms.ToTensor()(img))
return ZSSRDataset(pil, sr_factor)
def main():
SIZE_IMG = 1024; # 224
global args
args = parser.parse_args()
# create model
if args.arch.startswith('resnet50') :
model = model_defs.resnet50_oneway(num_classes=2);
elif args.arch.startswith('vgg16_bn') :
model = model_defs.vgg16_bn_oneway(num_classes=2, pretrained=True);
elif args.arch.startswith('segnet') :
model = model_defs.segnet(num_classes=2, pretrained=True);
model = nn.DataParallel(model.model); # dirty trick
#model = nn.DataParallel(model); # dirty trick
# open log file
if args.train == 1 :
log_dir = 'logs';
log_name = args.arch + '_new.csv';
if not os.path.isdir(log_dir) :
os.mkdir(log_dir);
log_handle = get_file_handle(os.path.join(log_dir, log_name), 'wb+');
log_handle.write('Epoch, LearningRate, Momentum, WeightDecay,' + \
'Loss, Precision, Recall, Accuracy(IoU), FgWeight, BgWeight\n');
log_handle.close();
# check model directory
model_dir = 'models';
if not os.path.isdir(model_dir) :
os.mkdir(model_dir);
# resume learning based on cmdline arguments
if ((args.start_epoch > 1) and (args.train==1)) :
load_epoch = args.start_epoch - 1;
elif (args.train==0) :
load_epoch = args.load_epoch;
else :
load_epoch = 0;
if load_epoch > 0 :
print("=> loading checkpoint for epoch = '{}'"
.format(load_epoch));
checkpoint_name = args.arch + '_ep_' + str(load_epoch) + '.pth.tar';
checkpoint = torch.load(os.path.join(model_dir, checkpoint_name));
model.load_state_dict(checkpoint['state_dict']);
# if args.arch.startswith('resnet50') :
# model = add_dropout2d(model);
# if args.arch.startswith('vgg16_bn') :
# model = add_dropout2d_vgg(model);
#model = add_conv_layers_vgg(model);
model.cuda(); # transfer to cuda
print(model);
mean = load_pickle('./mean');
std = load_pickle('./std');
if args.train == 1 :
train_data_dir, train_gt_dir = args.data, args.gt;
train_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(train_data_dir, train_gt_dir,
transform_joint=transforms.Compose_Joint([
transforms.RandomCrop(SIZE_IMG),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(), ]),
transform=transforms.Compose([
transforms.ColorJitter(0.3, 0.3, 0.3, 0),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),]),
target_transform=transforms.Compose([
transforms.ToTensorTarget(),]),
do_copy=True,
),
batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True);
weights = torch.from_numpy(np.array([1.,1.2])).float();
criterion = nn.CrossEntropyLoss(weights).cuda();
if args.optim == 'adam' :
optimizer = torch.optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay);
elif args.optim == 'sgd' :
optimizer = torch.optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay);
for epoch in range(args.start_epoch, args.end_epoch+1):
# train for one epoch
stats_epoch = train(train_loader, model, criterion, optimizer, epoch);
model_name = args.arch + '_ep_' + str(epoch) + '.pth.tar';
# get current parameters of optimizer
for param_group in optimizer.param_groups :
cur_lr = param_group['lr'];
cur_wd = param_group['weight_decay'];
if param_group.has_key('momentum') :
cur_momentum = param_group['momentum'];
else :
cur_momentum = 'n/a';
break; # constant parameters throughout the network
if epoch % args.save_interval == 0 :
state = {
'epoch': epoch,
'arch': args.arch,
'state_dict': model.state_dict(),
'learning_rate': cur_lr,
'moemntum': cur_momentum,
'weight_decay': cur_wd,
'fg_weight': weights[1],
'bg_weight': weights[0],
};
torch.save(state, os.path.join(model_dir, model_name));
# write logs using logHandle
log_handle = get_file_handle(os.path.join(log_dir, log_name), 'ab');
log_handle.write(str(epoch) + ',' +
str(cur_lr) + ',' +
str(cur_momentum) + ',' +
str(cur_wd) + ',' +
str(stats_epoch['loss']) + ',' +
str(stats_epoch['prec']) + ',' +
str(stats_epoch['recall']) + ',' +
str(stats_epoch['acc']) + ',' +
str(weights[1]) + ',' +
str(weights[0]) + '\n');
log_handle.close();
# adjust_learning_rate(optimizer, epoch, 10); # adjust learning rate
elif args.train == 0 : # test
testdir = args.data;
outdir = args.out;
# stride = args.test_stride;
# test_batch_size = args.test_batch_size;
test_transformer = transforms.Compose([
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),]);
# test(testdir, outdir, test_transformer, model, load_epoch, stride, SIZE_IMG);
# test_batch_form(testdir, outdir, test_transformer, model, load_epoch,
# stride, SIZE_IMG, test_batch_size);
test_full_res(testdir, outdir, test_transformer, model, load_epoch, SIZE_IMG);
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
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'))
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_dataset(name=args.dataset)
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_train_p = T.Compose([
T.Random2DTranslation(256, 128),
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]),
])
transform_test_p = T.Compose([
T.Resize((256, 128)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_file = 'data/cuhk_train.pkl'
test_file = 'data/cuhk_test.pkl'
gallery_file = 'data/cuhk_gallery.pkl'
data_root = args.data_root
dataset_train = CUHKGroup(train_file, data_root, True, transform_train,
transform_train_p)
dataset_test = CUHKGroup(test_file, data_root, False, transform_test,
transform_test_p)
dataset_query = CUHKGroup(test_file, data_root, False, transform_test,
transform_test_p)
dataset_gallery = CUHKGroup(gallery_file, data_root, False, transform_test,
transform_test_p)
pin_memory = True if use_gpu else False
if args.xent_only:
trainloader = DataLoader(
dataset_train,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
else:
trainloader = DataLoader(
dataset_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(
dataset_test,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
querygalleryloader = DataLoader(
dataset_query,
batch_size=args.gallery_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
galleryloader = DataLoader(
dataset_gallery,
batch_size=args.gallery_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
print("Initializing model: {}".format(args.arch))
if args.xent_only:
# model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent'})
model = models.init_model(name=args.arch,
num_classes=dataset_train.num_train_gids,
loss={'xent'})
else:
# model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent', 'htri'})
model = models.init_model(
name=args.arch,
num_classes=dataset_train.num_train_gids,
num_person_classes=dataset_train.num_train_pids,
loss={'xent', 'htri'})
#criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_gids, use_gpu=use_gpu)
#criterion_xent_person = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_pids, use_gpu=use_gpu)
if os.path.exists(args.pretrained_model):
print("Loading checkpoint from '{}'".format(args.pretrained_model))
checkpoint = torch.load(args.pretrained_model)
model_dict = model.state_dict()
pretrain_dict = checkpoint['state_dict']
pretrain_dict = {
k: v
for k, v in pretrain_dict.items() if k in model_dict
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
criterion_xent = nn.CrossEntropyLoss(ignore_index=-1)
criterion_xent_person = nn.CrossEntropyLoss(ignore_index=-1)
criterion_htri = TripletLoss(margin=args.margin)
criterion_pair = ContrastiveLoss(margin=args.margin)
criterion_htri_filter = TripletLossFilter(margin=args.margin)
criterion_permutation = PermutationLoss()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
if args.stepsize > 0:
if args.warmup:
scheduler = WarmupMultiStepLR(optimizer, [200, 400, 600])
else:
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.stepsize,
gamma=args.gamma)
start_epoch = args.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
test_gcn_person_batch(model, queryloader, querygalleryloader,
galleryloader, args.pool, use_gpu)
#test_gcn_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu)
#test_gcn(model, queryloader, galleryloader, args.pool, use_gpu)
#test(model, queryloader, galleryloader, args.pool, use_gpu)
return
start_time = time.time()
best_rank1 = -np.inf
for epoch in range(start_epoch, args.max_epoch):
#print("==> Epoch {}/{} lr:{}".format(epoch + 1, args.max_epoch, scheduler.get_lr()[0]))
train_gcn(model, criterion_xent, criterion_xent_person, criterion_pair,
criterion_htri_filter, criterion_htri, criterion_permutation,
optimizer, trainloader, use_gpu)
#train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu)
if args.stepsize > 0: scheduler.step()
if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (
epoch + 1) == args.max_epoch:
print("==> Test")
rank1 = test_gcn_person_batch(model, queryloader,
querygalleryloader, galleryloader,
args.pool, use_gpu)
#rank1 = test_gcn(model, queryloader, galleryloader, args.pool, use_gpu=False)
#rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu)
is_best = rank1 > best_rank1
if is_best: best_rank1 = rank1
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'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
