def main(args):
# Data
print('==> Preparing data..')
_size = 32
transform_train = transforms.Compose([
transforms.Resize(size=_size),
transforms.RandomResizedCrop(size=_size, scale=(0.2, 1.)),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomGrayscale(p=0.2),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.Resize(size=_size),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset = datasets.CIFAR10Instance(root='./data',
train=True,
download=True,
transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
testset = datasets.CIFAR10Instance(root='./data',
train=False,
download=True,
transform=transform_test)
testloader = torch.utils.data.DataLoader(testset,
batch_size=100,
shuffle=False,
num_workers=4)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
ndata = trainset.__len__()
print('==> Building model..')
net = models.__dict__['ResNet18'](low_dim=args.low_dim)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if device == 'cuda':
net = torch.nn.DataParallel(net,
device_ids=range(
torch.cuda.device_count()))
cudnn.benchmark = True
criterion = ICRcriterion()
# define loss function: inner product loss within each mini-batch
uel_criterion = BatchCriterion(args.batch_m, args.batch_t, args.batch_size,
ndata)
net.to(device)
criterion.to(device)
uel_criterion.to(device)
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
if args.test_only or len(args.resume) > 0:
# Load checkpoint.
model_path = 'checkpoint/' + args.resume
print('==> Resuming from checkpoint..')
assert os.path.isdir(
args.model_dir), 'Error: no checkpoint directory found!'
checkpoint = torch.load(model_path)
net.load_state_dict(checkpoint['net'])
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch']
# define leminiscate
if args.test_only and len(args.resume) > 0:
trainFeatures, feature_index = compute_feature(trainloader, net,
len(trainset), args)
lemniscate = LinearAverage(torch.tensor(trainFeatures), args.low_dim,
ndata, args.nce_t, args.nce_m)
else:
lemniscate = LinearAverage(torch.tensor([]), args.low_dim, ndata,
args.nce_t, args.nce_m)
lemniscate.to(device)
# define optimizer
optimizer = torch.optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
# optimizer2 = torch.optim.SGD(net2.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
# test acc
if args.test_only:
acc = kNN(0,
net,
trainloader,
testloader,
200,
args.batch_t,
ndata,
low_dim=args.low_dim)
exit(0)
if len(args.resume) > 0:
best_acc = best_acc
start_epoch = start_epoch + 1
else:
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
icr2 = ICRDiscovery(ndata)
# init_cluster_num = 20000
for round in range(5):
for epoch in range(start_epoch, 200):
#### get Features
# trainFeatures are trainloader features and shuffle=True, so feature_index is match data
trainFeatures, feature_index = compute_feature(
trainloader, net, len(trainset), args)
if round == 0:
y = -1 * math.log10(ndata) / 200 * epoch + math.log10(ndata)
cluster_num = int(math.pow(10, y))
if cluster_num <= args.nmb_cluster:
cluster_num = args.nmb_cluster
print('cluster number: ' + str(cluster_num))
###clustering algorithm to use
# faiss cluster
deepcluster = clustering.__dict__[args.clustering](
int(cluster_num))
#### Features to clustering
clustering_loss = deepcluster.cluster(trainFeatures,
feature_index,
verbose=args.verbose)
L = np.array(deepcluster.images_lists)
image_dict = deepcluster.images_dict
print('create ICR ...')
# icr = ICRDiscovery(ndata)
# if args.test_only and len(args.resume) > 0:
# icr = cluster_assign(icr, L, trainFeatures, feature_index, trainset,
# cluster_ratio + epoch*((1-cluster_ratio)/250))
icrtime = time.time()
# icr = cluster_assign(epoch, L, trainFeatures, feature_index, 1, 1)
if epoch < args.warm_epoch:
icr = cluster_assign(epoch, L, trainFeatures,
feature_index, args.cluster_ratio, 1)
else:
icr = PreScore(epoch, L, image_dict, trainFeatures,
feature_index, trainset, args.high_ratio,
args.cluster_ratio, args.alpha, args.beta)
print('calculate ICR time is: {}'.format(time.time() -
icrtime))
writer.add_scalar('icr_time', (time.time() - icrtime),
epoch + round * 200)
else:
cluster_num = args.nmb_cluster
print('cluster number: ' + str(cluster_num))
###clustering algorithm to use
# faiss cluster
deepcluster = clustering.__dict__[args.clustering](
int(cluster_num))
#### Features to clustering
clustering_loss = deepcluster.cluster(trainFeatures,
feature_index,
verbose=args.verbose)
L = np.array(deepcluster.images_lists)
image_dict = deepcluster.images_dict
print('create ICR ...')
# icr = ICRDiscovery(ndata)
# if args.test_only and len(args.resume) > 0:
# icr = cluster_assign(icr, L, trainFeatures, feature_index, trainset,
# cluster_ratio + epoch*((1-cluster_ratio)/250))
icrtime = time.time()
# icr = cluster_assign(epoch, L, trainFeatures, feature_index, 1, 1)
icr = PreScore(epoch, L, image_dict, trainFeatures,
feature_index, trainset, args.high_ratio,
args.cluster_ratio, args.alpha, args.beta)
print('calculate ICR time is: {}'.format(time.time() -
icrtime))
writer.add_scalar('icr_time', (time.time() - icrtime),
epoch + round * 200)
# else:
# icr = cluster_assign(icr, L, trainFeatures, feature_index, trainset, 0.2 + epoch*0.004)
# print(icr.neighbours)
icr2 = train(epoch, net, optimizer, lemniscate, criterion,
uel_criterion, trainloader, icr, icr2,
args.stage_update, args.lr, device, round)
print('----------Evaluation---------')
start = time.time()
acc = kNN(0,
net,
trainloader,
testloader,
200,
args.batch_t,
ndata,
low_dim=args.low_dim)
print("Evaluation Time: '{}'s".format(time.time() - start))
writer.add_scalar('nn_acc', acc, epoch + round * 200)
if acc > best_acc:
print('Saving..')
state = {
'net': net.state_dict(),
'acc': acc,
'epoch': epoch,
}
if not os.path.isdir(args.model_dir):
os.mkdir(args.model_dir)
torch.save(state,
'./checkpoint/ckpt_best_round_{}.t7'.format(round))
best_acc = acc
state = {
'net': net.state_dict(),
'acc': acc,
'epoch': epoch,
}
torch.save(state,
'./checkpoint/ckpt_last_round_{}.t7'.format(round))
print(
'[Round]: {} [Epoch]: {} \t accuracy: {}% \t (best acc: {}%)'.
format(round, epoch, acc, best_acc))
# define loss function
if hasattr(lemniscate, 'K'):
criterion = NCECriterion(ndata)
else:
criterion = nn.CrossEntropyLoss()
CE = nn.CrossEntropyLoss()
ID = InstanceDiscrimination(tau=1.0)
FD = FeatureDecorrelation(args.low_dim, tau2=2.0)
net.to(device)
lemniscate.to(device)
criterion.to(device)
ID.to(device)
FD.to(device)
LA.to(device)
CE.to(device)
if args.test_only:
acc = kNN(0, net, lemniscate, trainloader, testloader, 200, args.nce_t, 1)
sys.exit(0)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
def adjust_learning_rate(optimizer, epoch):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
lr = args.lr