query_cam, query_label, query_dict = get_id(query_path)
gallery_cam, gallery_label, gallery_dict = get_id(gallery_path)
"""
for key, value in query_dict.items():
print(key)
for key, value in gallery_dict.items():
print(key)
"""
print('-------test-----------')
if opt.loss_type == 'soft':
if opt.name == 'resnet_50':
model_structure = ft_net(751)
elif opt.name == 'resnext_50':
model_structure = resnext50(num_classes=751)
elif opt.name == 'densenet_121':
model_structure = ft_net_dense(751)
model = load_network(model_structure)
# Change to test mode
model = model.eval()
if use_gpu:
model = model.cuda()
for p in model.parameters():
p.requires_grad = False
# print(p.requires_grad)
gallery_feature_2048_adv = torch.FloatTensor()
gallery_feature_2048 = torch.FloatTensor()
torch.save(network.cpu().state_dict(), save_path)
if torch.cuda.is_available:
network.cuda(gpu_ids[0])
######################################################################
# Finetuning the convnet
# ----------------------
#
# Load a pretrainied model and reset final fully connected layer.
#
if opt.use_dense:
model = ft_net_dense(len(class_names))
elif opt.use_resnext:
model = resnext50(num_classes=len(class_names))
model.load_state_dict(torch.load('resnext_pretrain.pth'),False)
else:
model = ft_net(len(class_names))
print(model)
if use_gpu:
model = model.cuda()
criterion = nn.CrossEntropyLoss()
#criterion = AngleLoss()
#ignored_params = list(map(id, model.model.fc.parameters() )) + list(map(id, model.fc0.parameters() ))
#base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
# Observe that all parameters are being optimized
#optimizer_ft = optim.SGD([
def main():
global args
args = parser.parse_args()
print('config: wd', args.weight_decay, 'lr', args.lr, 'batch_size',
args.batch_size, 'num_gpus', args.num_gpus)
# torch.cuda.device_count(): 返回可得到的GPU数量
iteration_size = args.num_gpus // torch.cuda.device_count(
) # do multiple iterations
assert iteration_size >= 1
args.weight_decay = args.weight_decay * iteration_size # will cancel out with lr
args.lr = args.lr / iteration_size
args.batch_size = args.batch_size // iteration_size
# 对于只有一块GPU来说，参数没有改变
print('real: wd', args.weight_decay, 'lr', args.lr, 'batch_size',
args.batch_size, 'iteration_size', iteration_size)
# 分布式处理部分
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
print("=> creating model '{}'".format(args.arch))
from resnext import resnext50
model = resnext50(num_classes=80)
# from resnext import resnext50_elastic
# model = resnext50_elastic(num_classes=80)
# from resnext_MulTask_12 import resnext50_elastic
# model = resnext50_elastic(num_classes=80) # 模型 经过预训练 # opts["num_labels"] = 14
# count number of parameters
count = 0
params = list()
for n, p in model.named_parameters():
if '.ups.' not in n:
params.append(p)
count += np.prod(p.size())
print('Parameters:', count / 1000000.0) # 参数的数量
# count flops
model = add_flops_counting_methods(model)
model.eval()
image = torch.randn(1, 3, 224, 224)
model.start_flops_count()
model(image).sum()
model.stop_flops_count()
print("GFLOPs",
model.compute_average_flops_cost() / 1000000000.0) # FLOP的个数
# normal code
model = torch.nn.DataParallel(model).cuda()
# BCE损失函数
criterion = nn.BCEWithLogitsLoss().cuda()
# SGD优化策略
optimizer = torch.optim.SGD([
{
'params': iter(params),
'lr': args.lr
},
],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume) # 加载ckpt文件
resume = ('module.fc.bias' in checkpoint['state_dict'] and
checkpoint['state_dict']['module.fc.bias'].size() == model.module.fc.bias.size()) or \
('module.classifier.bias' in checkpoint['state_dict'] and
checkpoint['state_dict']['module.classifier.bias'].size() == model.module.classifier.bias.size())
if resume:
# True resume: resume training on MS-COCO # 在MS-COCO上 评估？
print()
print("resume training on MS-COCO...")
print("在MS-COCO上 评估...")
print()
model.load_state_dict(checkpoint['state_dict'], strict=False)
optimizer.load_state_dict(
checkpoint['optimizer']
) if 'optimizer' in checkpoint else print('no optimizer found')
args.start_epoch = checkpoint[
'epoch'] if 'epoch' in checkpoint else args.start_epoch
else:
# Fake resume: transfer from ImageNet # 从ImageNet——>MS-COCO 训练？
print()
print("transfer from ImageNet...")
print("从ImageNet——>MS-COCO 训练...")
print()
for n, p in list(checkpoint['state_dict'].items()):
if 'classifier' in n or 'fc' in n:
print(n, 'deleted from state_dict')
del checkpoint['state_dict'][n]
model.load_state_dict(checkpoint['state_dict'], strict=False)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume,
checkpoint['epoch'] if 'epoch' in checkpoint else 'unknown'))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# 提升一点训练速度，没什么额外开销，一般都会加
# 仅限于非多尺度训练！否则效果更差！
cudnn.benchmark = True
# Data loading code
# ms-coco正则化部分
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
# 训练集处理
train_dataset = CocoDetection(
os.path.join(args.data, 'train2014'),
os.path.join(args.data, 'annotations/instances_train2014.json'),
transforms.Compose([
transforms.RandomResizedCrop(224), # 随机裁剪
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ToTensor(),
normalize, # 正则化
]))
# print("train_dataset: ", train_dataset)
# 验证集处理
val_dataset = CocoDetection(
os.path.join(args.data, 'val2014'),
os.path.join(args.data, 'annotations/instances_val2014.json'),
transforms.Compose([
transforms.Resize((224, 224)), # 调整图像大小
transforms.ToTensor(),
normalize, # 正则化
]))
train_sampler = torch.utils.data.sampler.RandomSampler(
train_dataset) # 随机采样器
# torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=default_collate,
# pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None)
# 1.dataset(Dataset): (image, label)形式，数据读取接口（比如torchvision.datasets.ImageFolder）或者自定义的数据接口的输出，
# 该输出是torch.utils.data.Dataset类的对象(或者继承自该类的自定义类的对象)。
# 2.batch_size: 批训练数据量的大小，根据具体情况设置即可。（默认：1）
# 3.shuffle: 打乱数据，一般在训练数据中会采用。（默认：False）
# 4.sampler(Sampler, optional): 从数据集中提取样本的策略。如果指定，“shuffle”必须为false。一般默认即可。
# 5.num_workers，这个参数必须大于等于0，其他大于0的数表示通过多个进程来导入数据，可以加快数据导入速度。（默认：0）
# 6.pin_memory (bool, optional)：数据加载器将把张量复制到CUDA内存中，然后返回它们。也就是一个数据拷贝的问题。
# 7.drop_last (bool, optional): 如果数据集大小不能被批大小整除，则设置为“true”以除去最后一个未完成的批。如果“false”那么最后一批将更小。（默认：false）
# 8.timeout(numeric, optional)：设置数据读取超时，但超过这个时间还没读取到数据的话就会报错。（默认：0）
# 训练集
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler,
drop_last=True)
# 验证集
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate_multi(val_loader, model, criterion) # 在验证集上测试数据，return返回
return
for epoch in range(args.start_epoch, args.epochs):
# 学习率调整
coco_adjust_learning_rate(optimizer, epoch)
# 模型训练
train_multi(train_loader, model, criterion, optimizer, epoch,
iteration_size)
print("***********************************************")
print("模型训练完第 " + str(epoch + 1) + " 轮，下面进行验证集实验...")
print("***********************************************")
# evaluate on validation set
# 模型验证
validate_multi(val_loader, model, criterion)
# 模型保存的位置
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
},
False,
filename='3-MSCOCO--model-train-demo/' + 'coco_' + args.arch +
'_checkpoint_' + str(epoch + 1) + '.pth.tar')
def main():
global args, best_prec1
args = parser.parse_args()
# create model
model = resnext.resnext50(baseWidth=args.base_width,
cardinality=args.cardinality)
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'ILSVRC2012_img_train')
valdir = os.path.join(args.data, 'ILSVRC2012_img_val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomRotation(10),
transforms.RandomResizedCrop(224),
transforms.ColorJitter(
brightness=0.4,
contrast=0.4,
saturation=0.4,
),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
# train_loader = torch.utils.data.DataLoader(
# datasets.ImageFolder(traindir, transforms.Compose([
# transforms.RandomSizedCrop(224),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# normalize,
# ])),
# batch_size=args.batch_size, shuffle=True,
# num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': 'resnet',
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=4)
if args.model == 'resnet18':
cnn = ResNet18(num_classes=num_classes)
elif args.model == 'resnet50':
from resnet import resnet50
cnn = resnet50(num_classes)
print(cnn)
elif args.model == 'resnext50':
from resnext import resnext50
cnn = resnext50(num_classes)
print(cnn)
elif args.model == 'densenet121':
from densenet import densenet121
cnn = densenet121(num_classes)
print(cnn)
elif args.model == 'wideresnet':
if args.dataset == 'svhn':
cnn = WideResNet(depth=16,
num_classes=num_classes,
widen_factor=8,
dropRate=0.4)
else:
cnn = WideResNet(depth=28,
num_classes=num_classes,
widen_factor=10,