def loader(train_size, test_size, args):
if args.data.startswith('cifar'):
if args.data == 'cifar10':
dataloader = datasets.CIFAR10
else:
dataloader = datasets.CIFAR100
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
transforms.RandomErasing(probability = 0.5, sh = 0.4, r1 = 0.3, ),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
elif args.data == 'mnist':
dataloader = datasets.MNIST
transform_train = transforms.Compose([
# https://github.com/hwalsuklee/tensorflow-mnist-cnn/blob/master/mnist_data.py
#transforms.RandomAffine(translate=0.12),
transforms.RandomCrop(28, padding=4),
transforms.RandomRotation((-15, 15)),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
])
elif args.data == 'fmnist':
dataloader = datasets.FashionMNIST
transform_train = transforms.Compose([
transforms.RandomCrop(28, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
transforms.RandomErasing(probability=0.5, sh=0.4, r1=0.3, mean=[0.4914]),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)),
])
else:
exit('Unknown dataset')
if args.aug == 0:
transform_train = transforms.ToTensor()
transform_test = transforms.ToTensor()
trainset = dataloader('./data/' + args.data.upper(), train=True, download=True, transform=transform_train)
train_loader = data.DataLoader(trainset, batch_size=train_size, shuffle=True, num_workers=0) # num_workers=0 is crucial for seed
testset = dataloader(root='./data/' + args.data.upper(), train=False, download=False, transform=transform_test)
test_loader = data.DataLoader(testset, batch_size=test_size, shuffle=False, num_workers=0)
return train_loader, test_loader, dataloader
def fashion_mnist_aug_loader(train_size, test_size, args):
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, )),
])
trainset = datasets.FashionMNIST('./data/MNIST-FASHION',
train=True,
download=True,
transform=transform_train)
train_loader = data.DataLoader(trainset,
batch_size=train_size,
shuffle=True)
testset = datasets.FashionMNIST(root='./data/MNIST-FASHION',
train=False,
download=False,
transform=transform_test)
test_loader = data.DataLoader(testset, batch_size=test_size, shuffle=False)
return train_loader, test_loader
def get_mini_imagenet_data_loader(task, num_per_class=1, split='train',shuffle = False,train_query_argue=False,train_support_argue=False):
normalize = transforms.Normalize(mean=[0.92206, 0.92206, 0.92206], std=[0.08426, 0.08426, 0.08426])
Arguementation = T.Compose([
T.RandomCrop(84, padding=8),
T.RandomHorizontalFlip(),
T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
T.ToTensor(),
normalize,
T.RandomErasing(0.5)
])
Arguementation_support = transforms.Compose([
transforms.Resize(84),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVer
transforms.ToTensor(),
normalize
])
# 针对Support不做增强
if split == 'train' and train_support_argue == False:
dataset = MiniImagenet(task,split=split,transform=transforms.Compose([transforms.ToTensor(),normalize]))
sampler = ClassBalancedSamplerOld(num_per_class,task.num_classes, task.train_num,shuffle=shuffle)
elif split == 'train' and train_support_argue == True:
dataset = MiniImagenet(task,split=split,transform=Arguementation_support)
sampler = ClassBalancedSamplerOld(num_per_class,task.num_classes, task.train_num,shuffle=shuffle)
# 只针对训练中的query set做增强，但对训练中的support set不做增强
else:
if train_query_argue == False:
dataset = MiniImagenet(task,split=split,transform=transforms.Compose([transforms.ToTensor(),normalize]))
sampler = ClassBalancedSampler(task.num_classes, task.test_num,shuffle=shuffle)
else:
dataset = MiniImagenet(task,split=split,transform=Arguementation)
sampler = ClassBalancedSampler(task.num_classes, task.test_num,shuffle=shuffle)
loader = DataLoader(dataset, batch_size=num_per_class*task.num_classes, sampler=sampler)
return loader
def load_data_fashion_mnist(batch_size,
root='./dataset',
use_normalize=False,
mean=None,
std=None):
"""Download the fashion mnist dataset and then load into memory."""
if use_normalize:
normalize = transforms.Normalize((0.1307, ), (0.3081, ))
# normalize = transforms.Normalize(mean=[mean], std=[std])
# train_augs = transforms.Compose([transforms.RandomCrop(28, padding=2),
# transforms.RandomHorizontalFlip(),
# #transforms.RandomRotation(10),
# transforms.RandomAffine(degrees=0, translate=(0.05, 0.05)),
# transforms.ToTensor(),
# normalize,
# transforms.RandomErasing()])
train_augs = transforms.Compose([
transforms.RandomCrop(28, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
transforms.RandomErasing(probability=0,
sh=0.4,
r1=0.3,
mean=[0.4914]),
])
test_augs = transforms.Compose([transforms.ToTensor(), normalize])
else:
train_augs = transforms.Compose([transforms.ToTensor()])
test_augs = transforms.Compose([transforms.ToTensor()])
mnist_train = torchvision.datasets.FashionMNIST(root=root,
train=True,
download=True,
transform=train_augs)
mnist_test = torchvision.datasets.FashionMNIST(root=root,
train=False,
download=True,
transform=test_augs)
if sys.platform.startswith('win'):
num_workers = 0 # 0表示不用额外的进程来加速读取数据
else:
num_workers = 4
train_iter = torch.utils.data.DataLoader(mnist_train,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(mnist_test,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
return train_iter, test_iter
def __init__(
self,
*,
# RandomResizeAndCrop params
crop_size,
min_scale=-0.2,
max_scale=0.5,
stretch_prob=0.8,
# AsymmetricColorJitter params
brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.5 / 3.14,
# Random[H,V]Flip params
asymmetric_jitter_prob=0.2,
do_flip=True,
):
super().__init__()
transforms = [
T.PILToTensor(),
T.AsymmetricColorJitter(brightness=brightness,
contrast=contrast,
saturation=saturation,
hue=hue,
p=asymmetric_jitter_prob),
T.RandomResizeAndCrop(crop_size=crop_size,
min_scale=min_scale,
max_scale=max_scale,
stretch_prob=stretch_prob),
]
if do_flip:
transforms += [
T.RandomHorizontalFlip(p=0.5),
T.RandomVerticalFlip(p=0.1)
]
transforms += [
T.ConvertImageDtype(torch.float32),
T.Normalize(mean=0.5, std=0.5), # map [0, 1] into [-1, 1]
T.RandomErasing(max_erase=2),
T.MakeValidFlowMask(),
T.ValidateModelInput(),
]
self.transforms = T.Compose(transforms)
If a str of ‘random’, erasing each pixel with random values.
'''
transform_img = T.Compose([
T.Resize((256,128)),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop([256,128]),
T.ToTensor(),
#T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
#torchvision.transforms.RandomErasing(p=1, scale=(0.02, 0.4), ratio=(0.3, 3.33))
#torchvision.transforms.RandomErasing(p=1, scale=(0.02, 0.4), ratio=(0.3, 3.33), value=(0.4914, 0.4822, 0.4465))
#torchvision.transforms.RandomErasing(p=1, scale=(0.02, 0.4), ratio=(0.3, 3.33), value=1)
#torchvision.transforms.RandomErasing(p=1, scale=(0.02, 0.4), ratio=(0.3, 3.33), value=2)
#torchvision.transforms.RandomErasing(p=1, scale=(0.02, 0.4), ratio=(0.3, 3.33), value=12)
#torchvision.transforms.RandomErasing(p=1, scale=(0.02, 0.4), ratio=(0.3, 3.33), value='random')
T.RandomErasing(probability=0.5, sh=0.4, mean=(0.4914, 0.4822, 0.4465)),
])
if __name__ == '__main__':
pth = r'D:\pycharm\LR_reid\osnet\deep-person-reid-master\data\market1501\bounding_box_train\0002_c1s1_000451_03.jpg'
img, img_path = read_image(pth)
img_trans = transform_img(img)
print(type(img_trans),img_trans.dtype,img_trans.shape,sep='\t')
#plt_image(img_trans)
image = PIL.Image.fromarray(torch.clamp(img_trans * 255, min=0, max=255).byte().permute(1, 2, 0).cpu().numpy())
#image = torchvision.transforms.functional.to_pil_image(img_trans) # Equivalently way
plt_image(image)
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(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
transforms.RandomErasing(
probability=args.p,
sh=args.sh,
r1=args.r1,
),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
if args.dataset == 'cifar10':
dataloader = datasets.CIFAR10
num_classes = 10
else:
dataloader = datasets.CIFAR100
num_classes = 100
trainset = dataloader(root='./data',
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='./data',
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()
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 = 'cifar-10-' + 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))
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 train(img_path, label_path, model_path, net, batch_size, max_epoch,
checkpoint, lr, iscpu):
isgpu = not iscpu and torch.cuda.is_available()
if isgpu: cudnn.benchmark = True
rtime = time.time()
if net == 'FACE_NET':
model = FaceNet(Global.vecnum)
MCP = layer.MarginCosineProduct(Global.vecnum,
Global.class_out,
s=Global.mcp_size,
m=Global.mcp_dis)
if not isgpu:
model = model.cpu()
MCP = MCP.cpu()
#model = torch.nn.DataParallel(model)
torch.manual_seed(int(rtime))
else:
model = model.cuda()
#model = torch.nn.DataParallel(model).cuda()
torch.cuda.manual_seed(int(rtime))
MCP = MCP.cuda()
print(model)
kwargs = {'num_workers': 16, 'pin_memory': True} if isgpu else {}
transformList = []
transformList.append(transforms.RandomHorizontalFlip())
#transformList.append(transforms.RandomResizedCrop((Global.img_w,Global.img_h),scale=(0.07,1.0),ratio=(0.4,2.5)))
#transformList.append(transforms.ColorJitter(0.7,0.3,0.3,0.3))
transformList.append(transforms.Resize((Global.img_w, Global.img_h)))
transformList.append(
transforms.RandomCrop((Global.img_w, Global.img_h), 12))
#transformList.append(transforms.Resize((Global.img_w,Global.img_h)))
transformList.append(transforms.ToTensor())
transformList.append(transforms.Randomblack(mean=[0.0]))
transformList.append(transforms.RandomErasing(mean=[0.0]))
transformList.append(
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
transformSequence = transforms.Compose(transformList)
data = Img_Data(img_filepath=img_path,
label_filepath=label_path,
transform=transformSequence)
data_loader = DataLoader(dataset=data,
batch_size=batch_size,
shuffle=True,
**kwargs)
val_kwargs = {'num_workers': 1, 'pin_memory': True} if isgpu else {}
val_transformList = []
val_transformList.append(
transforms.Resize((Global.img_w, Global.img_h)))
val_transformList.append(transforms.ToTensor())
val_transformList.append(
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
val_transformSequence = transforms.Compose(val_transformList)
valdata = Img_Data(img_filepath=Global.val_img_path,
label_filepath=Global.val_label_path,
transform=val_transformSequence)
val_data_loader = DataLoader(dataset=valdata,
batch_size=1,
shuffle=False,
**val_kwargs)
loss = torch.nn.CrossEntropyLoss()
#optimizer = optim.SGD([{'params': model.parameters()}, {'params': MCP.parameters()}], lr=lr,momentum=0.9)
optimizer = optim.Adam([{
'params': model.parameters()
}, {
'params': MCP.parameters()
}],
lr=lr,
weight_decay=0,
amsgrad=True)
if checkpoint == True:
modelCheckpoint = torch.load(model_path)
model.load_state_dict(modelCheckpoint['state_dict'])
#optimizer.load_state_dict(modelCheckpoint['optimizer'])
MCP.load_state_dict(modelCheckpoint['MCP'])
start = modelCheckpoint['epoch'] + 1
else:
start = 0
for epoch_id in range(start, max_epoch):
oldtime = time.time()
loss_train = Train_Proc.epoch_train(model, data_loader, optimizer,
loss, MCP, isgpu)
newtime = time.time()
torch.save(
{
'epoch': epoch_id,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'MCP': MCP.state_dict()
}, './models/' + 'model-' + str(epoch_id) + '.pth.tar')
totaltime = newtime - oldtime
val_sum1, val_sum2 = Train_Proc.validate(model, val_data_loader,
isgpu)
logging.info(
'Epoch %d : training_loss = %.6f time = %d val1 = %.6f val2 = %.6f %.6f'
% (epoch_id, loss_train, totaltime, val_sum1, val_sum2,
val_sum1 - val_sum2))
with open(Global.loss_log, 'a+') as ftxt:
ftxt.write(
'\nEpoch %d : training_loss = %.6f time = %d val1 = %.6f val2 = %.6f %.6f\n'
% (epoch_id, loss_train, totaltime, val_sum1, val_sum2,
val_sum1 - val_sum2))
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)
############################# Hyper-parameters ################################
alpha = 1.0
beta = 1.0
at_margin = 1
pixel_mean = [0.485, 0.456, 0.406]
pixel_std = [0.229, 0.224, 0.225]
inp_size = [384, 128]
# transforms
transforms_list = transforms.Compose([
transforms.RectScale(*inp_size),
transforms.RandomHorizontalFlip(),
transforms.Pad(10),
transforms.RandomCrop(inp_size),
transforms.ToTensor(),
transforms.Normalize(mean=pixel_mean, std=pixel_std),
transforms.RandomErasing(probability=0.5, mean=pixel_mean)
])
test_transforms_list = transforms.Compose([
transforms.RectScale(*inp_size),
transforms.ToTensor(),
transforms.Normalize(mean=pixel_mean, std=pixel_std)
])
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]),
T.RandomErasing(probability=0.5, 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.xent_only:
trainloader = DataLoader(
VideoDataset(dataset.train,
seq_len=args.seq_len,
sample='random',
transform=transform_train),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
else:
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))
if args.arch == 'resnet503d':
model = resnet3d.resnet50(num_classes=dataset.num_train_pids,
sample_width=args.width,
sample_height=args.height,
sample_duration=args.seq_len)
if not os.path.exists(args.pretrained_model):
raise IOError("Can't find pretrained model: {}".format(
args.pretrained_model))
print("Loading checkpoint from '{}'".format(args.pretrained_model))
checkpoint = torch.load(args.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:
#model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, dropout=args.dropout, nhid=args.nhid, nheads=args.nheads, concat=args.concat, loss={'xent', 'htri'})
model = models.init_model(name=args.arch,
pool_size=8,
input_shape=2048,
n_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
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 = 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)
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 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'] + 1
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()
best_rank1 = -np.inf
if args.arch == 'resnet503d':
torch.backends.cudnn.benchmark = False
'''
adj1 = build_adj_full_full(4, args.part1)
adj2 = build_adj_full_full(4, args.part2)
adj3 = build_adj_full_full(4, args.part3)
if use_gpu:
adj1 = adj1.cuda()
adj2 = adj2.cuda()
adj2 = adj2.cuda()
adj1 = Variable(adj1)
adj2 = Variable(adj2)
adj3 = Variable(adj3)
'''
torch.cuda.empty_cache()
for epoch in range(start_epoch, args.max_epoch):
print("==> Epoch {}/{} lr:{}".format(epoch + 1, args.max_epoch,
scheduler.get_lr()[0]))
#train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu, adj1, adj2, adj3)
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(model, queryloader, galleryloader, args.pool, use_gpu, adj1, adj2, adj3)
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'))
torch.cuda.empty_cache()
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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)
# transforms.Resize((opt.imageSize, opt.imageSize)),
# transforms.RandomCrop(opt.imageSize),
# transforms.RandomHorizontalFlip(),
# transforms.ColorJitter(brightness=.1, contrast=.1, saturation=.1, hue=.1),
# transforms.ToTensor(),
# transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
# ])
transform_train = transforms.Compose([
transforms.Resize((84, 84), interpolation=3),
transforms.RandomCrop(84, padding=8),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
transforms.RandomErasing(0.5)
])
transform_test = transforms.Compose([
transforms.Resize((84, 84), interpolation=3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
trainset = Image_dataset(
data_dir='/home/lemon/few-shot/DN4/dataset/miniImageNet/mini-imagenet',
mode='train',
image_size=opt.imageSize,
transform=transform_train,
episode_num=opt.episode_train_num,
def main():
runId = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, runId)
if not os.path.exists(cfg.OUTPUT_DIR):
os.mkdir(cfg.OUTPUT_DIR)
print(cfg.OUTPUT_DIR)
torch.manual_seed(cfg.RANDOM_SEED)
random.seed(cfg.RANDOM_SEED)
np.random.seed(cfg.RANDOM_SEED)
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
use_gpu = torch.cuda.is_available() and cfg.MODEL.DEVICE == "cuda"
if not cfg.EVALUATE_ONLY:
sys.stdout = Logger(osp.join(cfg.OUTPUT_DIR, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(cfg.OUTPUT_DIR, 'log_test.txt'))
print("==========\nConfigs:{}\n==========".format(cfg))
if use_gpu:
print("Currently using GPU {}".format(cfg.MODEL.DEVICE_ID))
cudnn.benchmark = True
torch.cuda.manual_seed_all(cfg.RANDOM_SEED)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(cfg.DATASETS.NAME))
dataset = data_manager.init_dataset(root=cfg.DATASETS.ROOT_DIR,
name=cfg.DATASETS.NAME)
print("Initializing model: {}".format(cfg.MODEL.NAME))
if cfg.MODEL.ARCH == 'video_baseline':
torch.backends.cudnn.benchmark = False
model = models.init_model(name=cfg.MODEL.ARCH,
num_classes=625,
pretrain_choice=cfg.MODEL.PRETRAIN_CHOICE,
last_stride=cfg.MODEL.LAST_STRIDE,
neck=cfg.MODEL.NECK,
model_name=cfg.MODEL.NAME,
neck_feat=cfg.TEST.NECK_FEAT,
model_path=cfg.MODEL.PRETRAIN_PATH)
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
transform_train = T.Compose([
T.Resize(cfg.INPUT.SIZE_TRAIN),
T.RandomHorizontalFlip(p=cfg.INPUT.PROB),
T.Pad(cfg.INPUT.PADDING),
T.RandomCrop(cfg.INPUT.SIZE_TRAIN),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
T.RandomErasing(probability=cfg.INPUT.RE_PROB,
mean=cfg.INPUT.PIXEL_MEAN)
])
transform_test = T.Compose([
T.Resize(cfg.INPUT.SIZE_TEST),
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
cfg.DATALOADER.NUM_WORKERS = 0
trainloader = DataLoader(VideoDataset(
dataset.train,
seq_len=cfg.DATASETS.SEQ_LEN,
sample=cfg.DATASETS.TRAIN_SAMPLE_METHOD,
transform=transform_train,
dataset_name=cfg.DATASETS.NAME),
sampler=RandomIdentitySampler(
dataset.train,
num_instances=cfg.DATALOADER.NUM_INSTANCE),
batch_size=cfg.SOLVER.SEQS_PER_BATCH,
num_workers=cfg.DATALOADER.NUM_WORKERS,
pin_memory=pin_memory,
drop_last=True)
queryloader = DataLoader(VideoDataset(
dataset.query,
seq_len=cfg.DATASETS.SEQ_LEN,
sample=cfg.DATASETS.TEST_SAMPLE_METHOD,
transform=transform_test,
max_seq_len=cfg.DATASETS.TEST_MAX_SEQ_NUM,
dataset_name=cfg.DATASETS.NAME),
batch_size=cfg.TEST.SEQS_PER_BATCH,
shuffle=False,
num_workers=cfg.DATALOADER.NUM_WORKERS,
pin_memory=pin_memory,
drop_last=False)
galleryloader = DataLoader(
VideoDataset(dataset.gallery,
seq_len=cfg.DATASETS.SEQ_LEN,
sample=cfg.DATASETS.TEST_SAMPLE_METHOD,
transform=transform_test,
max_seq_len=cfg.DATASETS.TEST_MAX_SEQ_NUM,
dataset_name=cfg.DATASETS.NAME),
batch_size=cfg.TEST.SEQS_PER_BATCH,
shuffle=False,
num_workers=cfg.DATALOADER.NUM_WORKERS,
pin_memory=pin_memory,
drop_last=False,
)
if cfg.MODEL.SYN_BN:
if use_gpu:
model = nn.DataParallel(model)
if cfg.SOLVER.FP_16:
model = apex.parallel.convert_syncbn_model(model)
model.cuda()
start_time = time.time()
xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids)
tent = TripletLoss(cfg.SOLVER.MARGIN)
optimizer = make_optimizer(cfg, model)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
# metrics = test(model, queryloader, galleryloader, cfg.TEST.TEMPORAL_POOL_METHOD, use_gpu)
no_rise = 0
best_rank1 = 0
start_epoch = 0
for epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCHS):
# if no_rise == 10:
# break
scheduler.step()
print("noriase:", no_rise)
print("==> Epoch {}/{}".format(epoch + 1, cfg.SOLVER.MAX_EPOCHS))
print("current lr:", scheduler.get_lr()[0])
train(model, trainloader, xent, tent, optimizer, use_gpu)
if cfg.SOLVER.EVAL_PERIOD > 0 and (
(epoch + 1) % cfg.SOLVER.EVAL_PERIOD == 0 or
(epoch + 1) == cfg.SOLVER.MAX_EPOCHS):
print("==> Test")
metrics = test(model, queryloader, galleryloader,
cfg.TEST.TEMPORAL_POOL_METHOD, use_gpu)
rank1 = metrics[0]
if rank1 > best_rank1:
best_rank1 = rank1
no_rise = 0
else:
no_rise += 1
continue
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
torch.save(
state_dict,
osp.join(
cfg.OUTPUT_DIR, "rank1_" + str(rank1) + '_checkpoint_ep' +
str(epoch + 1) + '.pth'))
# best_p = osp.join(cfg.OUTPUT_DIR, "rank1_" + str(rank1) + '_checkpoint_ep' + str(epoch + 1) + '.pth')
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.checkpoint):
os.makedirs(args.checkpoint)
# Data
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
transforms.RandomErasing(probability = args.p, sh = args.sh, r1 = args.r1, ),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.train_batch, shuffle=True, num_workers=4)
testset = datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=args.test_batch, shuffle=False, num_workers=4)
def __init__(self, args, use_gpu):
super(DataManager, self).__init__()
self.args = args
self.use_gpu = use_gpu
print("Initializing dataset {}".format(args.dataset))
dataset = datasets.init_imgfewshot_dataset(name=args.dataset)
if args.load:
transform_train = T.Compose([
T.RandomCrop(84, padding=8),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
T.RandomErasing(0.5)
])
transform_test = T.Compose([
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
else:
transform_train = T.Compose([
T.Resize((args.height, args.width), interpolation=3),
T.RandomCrop(args.height, padding=8),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
T.RandomErasing(0.5)
])
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]),
])
pin_memory = True if use_gpu else False
self.trainloader = DataLoader(
dataset_loader.init_loader(
name='train_loader',
dataset=dataset.train,
labels2inds=dataset.train_labels2inds,
labelIds=dataset.train_labelIds,
nKnovel=args.nKnovel,
nExemplars=args.nExemplars,
nTestNovel=args.train_nTestNovel,
epoch_size=args.train_epoch_size,
transform=transform_train,
load=args.load,
),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
self.valloader = DataLoader(
dataset_loader.init_loader(
name='test_loader',
dataset=dataset.val,
labels2inds=dataset.val_labels2inds,
labelIds=dataset.val_labelIds,
nKnovel=args.nKnovel,
nExemplars=args.nExemplars,
nTestNovel=args.nTestNovel,
epoch_size=args.epoch_size,
transform=transform_test,
load=args.load,
),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
self.testloader = DataLoader(
dataset_loader.init_loader(
name='test_loader',
dataset=dataset.test,
labels2inds=dataset.test_labels2inds,
labelIds=dataset.test_labelIds,
nKnovel=args.nKnovel,
nExemplars=args.nExemplars,
nTestNovel=args.nTestNovel,
epoch_size=args.epoch_size,
transform=transform_test,
load=args.load,
),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
def main():
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))
use_gpu = torch.cuda.is_available()
np.random.seed(args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.deterministic = True
cudnn.benchmark = True
print("Initializing train dataset {}".format(args.train_dataset))
train_dataset = data_manager.init_dataset(name=args.train_dataset)
print("Initializing test dataset {}".format(args.test_dataset))
test_dataset = data_manager.init_dataset(name=args.test_dataset)
# print("Initializing train dataset {}".format(args.train_dataset, split_id=6))
# train_dataset = data_manager.init_dataset(name=args.train_dataset)
# print("Initializing test dataset {}".format(args.test_dataset, split_id=6))
# test_dataset = data_manager.init_dataset(name=args.test_dataset)
transform_train = T.Compose([
T.Resize([args.height, args.width]),
T.RandomHorizontalFlip(),
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]),
T.RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406])
])
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
# random_snip first_snip constrain_random evenly
trainloader = DataLoader(
VideoDataset(train_dataset.train,
seq_len=args.seq_len,
sample='constrain_random',
transform=transform_train),
sampler=RandomIdentitySampler(train_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(test_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(test_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=train_dataset.num_train_pids,
loss={'xent', 'htri'})
print("Model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1000000.0))
print("load model {0} from {1}".format(args.arch, args.load_model))
if args.load_model != '':
pretrained_model = torch.load(args.load_model)
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_model['state_dict'].items()
if k in model_dict
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
start_epoch = pretrained_model['epoch'] + 1
best_rank1 = pretrained_model['rank1']
else:
start_epoch = args.start_epoch
best_rank1 = -np.inf
criterion = dict()
criterion['triplet'] = WeightedRegularizedTriplet()
criterion['xent'] = CrossEntropyLabelSmooth(
num_classes=train_dataset.num_train_pids)
criterion['center'] = CenterLoss(num_classes=train_dataset.num_train_pids,
feat_dim=512,
use_gpu=True)
print(criterion)
optimizer = dict()
optimizer['model'] = model.get_optimizer(args)
optimizer['center'] = torch.optim.SGD(criterion['center'].parameters(),
lr=0.5)
scheduler = lr_scheduler.MultiStepLR(optimizer['model'],
milestones=args.stepsize,
gamma=args.gamma)
print(model)
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
distmat = test(model,
queryloader,
galleryloader,
args.pool,
use_gpu,
return_distmat=True)
return
start_time = time.time()
train_time = 0
best_epoch = args.start_epoch
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
scheduler.step()
print('Epoch', epoch, 'lr', scheduler.get_lr()[0])
start_train_time = time.time()
train(epoch, model, criterion, optimizer, trainloader, use_gpu)
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(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 __init__(self, args, use_gpu):
super(DataManager, self).__init__()
self.args = args
self.use_gpu = use_gpu
print("Initializing dataset {}".format(args.dataset))
dataset = datasets.init_imgfewshot_dataset(name=args.dataset,
root=args.root)
if args.load:
transform_train = T.Compose([
T.RandomCrop(84, padding=8),
T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
T.RandomErasing(0.5)
])
transform_test = T.Compose([
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
elif args.cifar:
transform_train = T.Compose([
T.RandomCrop(32, padding=4),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.507, 0.487, 0.441],
std=[0.267, 0.256, 0.276]),
T.RandomErasing(0.5)
])
transform_test = T.Compose([
T.ToTensor(),
T.Normalize(mean=[0.507, 0.487, 0.441],
std=[0.267, 0.256, 0.276]),
])
else:
transform_train = T.Compose([
T.RandomResizedCrop((84, 84)),
T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
T.RandomErasing(0.5)
])
transform_test = T.Compose([
T.Resize((92, 92)),
T.CenterCrop((84, 84)),
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
self.trainloader = DataLoader(
dataset_loader.init_loader(
name='train_loader',
dataset=dataset.train,
labels2inds=dataset.train_labels2inds,
labelIds=dataset.train_labelIds,
nKnovel=args.nKnovel,
nExemplars=args.nExemplars,
nTestNovel=args.train_nTestNovel,
epoch_size=args.train_epoch_size,
transform=transform_train,
load=args.load,
tiered=args.tiered,
),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
self.valloader = DataLoader(
dataset_loader.init_loader(
name='test_loader',
dataset=dataset.val,
labels2inds=dataset.val_labels2inds,
labelIds=dataset.val_labelIds,
nKnovel=args.nKnovel,
nExemplars=args.nExemplars,
nTestNovel=args.nTestNovel,
epoch_size=args.epoch_size,
transform=transform_test,
load=args.load,
tiered=args.tiered,
),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
self.testloader = DataLoader(
dataset_loader.init_loader(
name='test_loader',
dataset=dataset.test,
labels2inds=dataset.test_labels2inds,
labelIds=dataset.test_labelIds,
nKnovel=args.nKnovel,
nExemplars=args.nExemplars,
nTestNovel=args.nTestNovel,
epoch_size=args.epoch_size,
transform=transform_test,
load=args.load,
tiered=args.tiered,
),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
