def plot(ds, d_name, inds=[[1, 2, 3], [1, 4, 5]], add=0):
d = os.path.join('output', d_name)
check_dir_exists([d])
for i in range(len(inds)):
imgs = []
for j in range(len(inds[i])):
tmp = 0 if j == 0 else add
imgs.append(ds[inds[i][j] + tmp][0])
save_image(imgs, os.path.join(d, 'img_%d.jpg' % (i)))
def generate_feature(output_file):
'''
Generate feature of given dataset.
1. laod model
2. for each data in train_loader => (encoded, (label, img_name))
3. output a dictionary {'features': [encoded], 'labels': [(label, img_name)]}
:param output_file: output file name
'''
model_name = 'model/%s_%s%s_model-%s.pkl' % (args.main_model, args.model,
'' if args.fea_c is None else
args.fea_c, args.dataset)
print(model_name)
assert os.path.exists(model_name)
print('Loading model ...')
if cuda:
mol = torch.load(model_name).to(device)
else:
mol = torch.load(model_name, map_location='cpu').to(device)
# print(torchsummary.summary(autoencoder, input_size=(3, HEIGHT, WEIGHT)))
train_loader = getDataLoader(args, kwargs)
check_dir_exists(['feature'])
features = []
labels = []
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
label = [(y[0][i], y[1][i]) for i in range(len(y[0]))]
labels.extend(label)
if args.main_model == 'AE':
feature, _ = mol(b_x)
elif args.main_model == 'VAE':
_, mu, std = mol(b_x)
feature = torch.cat([mu, std])
elif args.main_model == 'AEClass':
feature, _, _ = mol(b_x)
elif args.main_model == 'vgg_classifier':
feature = mol.get_encode_features(b_x)
f = feature.cpu() if cuda else feature
f = f.data.view(b_x.shape[0], -1).numpy().tolist()
features.extend(f)
if step % 10 == 0:
print('Step %d finished.' % step)
out = {'features': features, 'labels': labels}
with open(output_file, 'w') as f:
json.dump(out, f)
return out
def get_features(mol_short='inception_v3'):
################################################################
# Arguments
################################################################
ae_args = train_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
start_time = time.time()
args = ae_args
model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model,
'' if args.fea_c is None else
args.fea_c, args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c,
args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
mol = inception_v3_features(pretrained=True, training=False).to(device)
# train_loader = getDataLoader(args, kwargs)
test_loader = getDataLoader(args, kwargs, train='test')
check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir])
t_per_img = []
for epoch in range(1):
step_time = time.time()
for step, (x, y) in enumerate(test_loader):
b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x)
t0 = time.time()
features = mol(b_x)
t_tmp = (time.time() - t0) / len(b_x) * 1000
t_per_img.append(t_tmp)
print(
'cost %.6fms per image this batch. cost %.6fms per image till now.'
% (t_tmp, np.mean(sorted(t_per_img)[1:-1])))
print('Finished. Totally cost %.2f' % (time.time() - start_time))
def __init__(self, name='inception_finetune_experiments'):
self.args = train_args()
self._get_gpu()
self.cuda = torch.cuda.is_available() and self.args.gpu != -2
self.device = torch.device("cuda" if self.cuda else "cpu")
self.kwargs = {'num_workers': 6, 'pin_memory': True}
self.mol_dir = 'model'
self.log_dir = os.path.join('log')
self.log_file = os.path.join(self.log_dir, 'log_%s.log' % name)
self.summary_dir = os.path.join('summary', self.args.name)
check_dir_exists([
self.mol_dir, self.log_dir,
os.path.dirname(self.summary_dir), self.summary_dir
])
self.writer = SummaryWriter(self.summary_dir)
logging.basicConfig(filename=self.log_file, level=logging.INFO)
def get_features(mol_short='resnet50'):
################################################################
# Arguments
################################################################
ae_args = train_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
start_time = time.time()
args = ae_args
model_name = 'model/%s_%s%s_model-%s.pkl' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
mol = ResNet().to(device)
test_loader = getDataLoader(args, kwargs, train='test')
loss_class = nn.CrossEntropyLoss().cuda(cuda)
check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir])
total, correct, top5correct, loss_total = 0, 0, 0, 0
t_per_img = []
for epoch in range(1):
step_time = time.time()
for step, (x, y) in enumerate(test_loader):
t0 = time.time()
b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x)
prob_class = mol.get_fc_features(b_x, return_both=False)
t_tmp = (time.time() - t0) / len(b_x) * 1000
t_per_img.append(t_tmp)
print('cost %.6fms per image this batch. cost %.6fms per image till now.' % (
t_tmp, np.mean(sorted(t_per_img)[1:-1])))
def train_cls(args):
################################################################
# Arguments
################################################################
cuda = args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
################################################################
# Model name & Directory
################################################################
mol_short = get_mol_short(args.model) + '(cls)'
model_name = 'model/%s_%s_model.pkl' % (
mol_short,
args.model,
)
print('[Model] model name is', model_name)
pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model,
'' if args.fea_c is None else args.fea_c,
args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c,
args.dataset)
log_dir = 'log/log_%s_%s%s_model-%s/' % \
(mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
remove_dir_exists([log_dir])
check_dir_exists([
'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir
])
writer = SummaryWriter(log_dir)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
mol = ModelCls(args).to(device)
################################################################
# Prepare Data & Optimizer
################################################################
print('Prepare data loader ...')
test_loader = getDataLoader(args, kwargs, train='test')
train_loader = getDataLoader(args, kwargs, train='train')
cover_val_loader = getDataLoader(args, kwargs, train='cover_validation')
cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample')
optimizer = torch.optim.Adam(mol.classifier.parameters(), lr=args.lr)
loss_class = nn.CrossEntropyLoss(reduction='none').cuda(cuda)
total, correct, top5correct, cnt = 0, 0, 0, 0
print('Start training ...')
for epoch in range(args.epoch):
################################################################
# Testing
################################################################
if epoch > 0:
mol.eval()
# Testing on ImageNet val
print('######### Testing on ImageNet val Dataset ###########')
test_loss_cls, test_acc, test_top5acc = test_cls(
test_loader, mol, cuda, 'Full')
writer.add_scalar('test_imagenet/loss_classifier', test_loss_cls,
epoch)
writer.add_scalar('test_imagenet/accuracy', test_acc, epoch)
writer.add_scalar('test_imagenet/top5accuracy', test_top5acc,
epoch)
if epoch % 5 == 4 or epoch == args.epoch - 1:
mol.eval()
evaluate_labeled_data(test_loader, mol, cuda)
################################################################
# Training
################################################################
step_time = time.time()
mol.train()
print('######### Training with %d batches total ##########' %
len(train_loader))
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
label = Variable(
torch.Tensor([y[2][i] for i in range(len(y[0]))]).long())
label = label.cuda() if cuda else label
prob_class = mol(b_x)
loss = loss_class(prob_class, label)
writer.add_scalar('train/loss_classifier', torch.mean(loss.data),
cnt)
optimizer.zero_grad()
torch.mean(loss).backward()
optimizer.step()
_, predicted = torch.max(prob_class.data, 1)
total += label.size(0)
correct += (predicted == label).sum().item()
top5pre = prob_class.topk(5, 1, True, True)
top5pre = top5pre[1].t()
top5correct += top5pre.eq(label.view(
1, -1).expand_as(top5pre)).sum().item()
writer.add_scalar('train/accuracy', correct / total, cnt)
writer.add_scalar('train/top5_accuracy', top5correct / total, cnt)
if step % 10 == 0:
if os.path.exists(model_name):
shutil.copy2(model_name,
model_name.split('.pkl')[0] + '_back.pkl')
torch.save(mol, model_name)
total_tmp = total if total != 0 else 1
print(
'[Training] Epoch:', epoch, 'Step:', step, '|',
'Time cost %.2f s; Classification error %.6f; '
'Accuracy %.3f%%; Top5 Accuracy %.3f%%' %
(time.time() - step_time, torch.mean(loss.data),
correct * 100 / total_tmp, top5correct * 100 / total_tmp))
correct, total, top5correct = 0, 0, 0
step_time = time.time()
cnt += 1
pass
def train():
################################################################
# Arguments
################################################################
vae_args = train_args()
cuda = vae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if vae_args.cuda else {}
start_time = time.time()
args = vae_args
model_name = 'model/VAE_%s%s_model-%s.pkl' % (
args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
pic_dir = 'res/VAE_%s%s-%s/' % (args.model, '' if args.fea_c is None else
args.fea_c, args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
if cuda:
vae = torch.load(model_name).to(device)
else:
vae = torch.load(model_name, map_location='cpu')
else:
vae = init_model(args.model).to(device)
train_loader = getDataLoader(args, kwargs)
optimizer = torch.optim.Adam(list(vae.parameters()), lr=args.lr)
check_dir_exists(['res/', 'model', pic_dir])
loss_val = None
for epoch in range(args.epoch):
step_time = time.time()
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
b_y = b_x.detach().cuda() if cuda else b_x.detach(
) # batch y, shape (batch, 32*32*3)
decoded, mu, std = vae(b_x)
if step % 100 == 0:
img_to_save = torch.cat([b_x.data, decoded.data])
save_image(img_to_save,
'%s/%s-%s.jpg' % (pic_dir, epoch, step))
# io.imsave('.xxx.jpg',img_to_save[0])
loss, bce, kld = loss_function(decoded, b_y, mu,
std) # mean square error
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # backpropagation, compute gradients
optimizer.step()
loss_val = 0.99 * loss_val + 0.01 * loss.data[
0] if loss_val is not None else loss.data[0]
if step % 10 == 0:
torch.save(vae, model_name)
print(
'Epoch:', epoch, 'Step:', step, '|',
'train loss %.6f; KLD %.6f; BCE %.6f; Time cost %.2f s' %
(loss_val, kld, bce, time.time() - step_time))
step_time = time.time()
print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(mol_short='inception_v3'):
################################################################
# Arguments
################################################################
ae_args = train_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
start_time = time.time()
args = ae_args
model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model,
'' if args.fea_c is None else
args.fea_c, args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c,
args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
mol = inception_v3(pretrained=True, training=False).to(device)
# mol.eval()
# train_loader = getDataLoader(args, kwargs)
test_loader = getDataLoader(args, kwargs, train='test')
# cover_loader = getDataLoader(args, kwargs, train='cover')
cover_val_loader = getDataLoader(args, kwargs, train='cover_validation')
cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample')
loss_class = nn.CrossEntropyLoss().cuda(cuda)
check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir])
# Evaluation
# check_dir_exists([os.path.join(evaluation_dir, 'cos'), os.path.join(evaluation_dir, 'distance')])
evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda,
evaluation_dir, args)
# encode_accuracy, encode_top5accuracy, fc_accuracy, fc_top5accuracy = evaluate_labeled_data(test_loader, mol, cuda)
# print('Encode accuracy:', np.mean(encode_accuracy))
# print('Encode top5 accuracy:', np.mean(encode_top5accuracy))
# print('Fc accuracy:', np.mean(fc_accuracy))
# print('Fc top5 accuracy:', np.mean(fc_top5accuracy))
total, correct, top5correct, loss_total = 0, 0, 0, 0
for epoch in range(1):
step_time = time.time()
for step, (x, y) in enumerate(test_loader):
b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x)
label = Variable(
torch.Tensor([y[2][i] for i in range(len(y[0]))]).long())
label = label.cuda() if cuda else label
prob_class = mol(b_x)
loss = loss_class(prob_class, label) # mean square error
_, predicted = torch.max(prob_class.data, 1)
total += label.size(0)
correct += (predicted == label).sum().item()
top5pre = prob_class.topk(5, 1, True, True)
top5pre = top5pre[1].t()
top5correct += top5pre.eq(label.view(
1, -1).expand_as(top5pre)).sum().item()
loss_total += loss.data[0] * label.size(0)
if step % 10 == 0:
torch.save(mol, model_name)
print(
'Epoch:', epoch, 'Step:', step, '|',
'test loss %.6f; Time cost %.2f s; Classification error %.6f; '
'Top1 Accuracy %.3f; Top5 Accuracy %.3f' %
(loss_total / total, time.time() - step_time, loss,
correct * 100 / total, top5correct * 100 / total))
step_time = time.time()
print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(mol_short='VGGClass', main_model=VGGClass):
################################################################
# Arguments
################################################################
ae_args = train_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
args = ae_args
mol_short = mol_short if args.name == '' else mol_short + '_' + args.name
log_dir = 'log/log_%s_%s%s_model-%s/' %\
(mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
writer = SummaryWriter(log_dir)
start_time = time.time()
model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model,
'' if args.fea_c is None else
args.fea_c, args.dataset)
print('[Model] model name is', model_name)
pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model,
'' if args.fea_c is None else args.fea_c,
args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c,
args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
if cuda:
mol = torch.load(model_name).to(device)
else:
mol = torch.load(model_name, map_location='cpu')
else:
print('Init model ...')
mol = main_model(args.fea_c).to(device)
print('Prepare data loader ...')
train_loader = getDataLoader(args, kwargs, p=args.imgnet_p)
test_loader = getDataLoader(args, kwargs, train='test')
# small_test_loader = getDataLoader(args, kwargs, train=False, p=10)
# cover_loader = getDataLoader(args, kwargs, train='cover')
cover_val_loader = getDataLoader(args, kwargs, train='cover_validation')
cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample')
# Optimizer & Loss function
if args.fea_c == 512:
optimizer1 = torch.optim.Adam(list(mol.classification.parameters()),
lr=args.lr)
else:
optimizer1 = torch.optim.Adam(list(mol.classification.parameters()) +
list(mol.small_features.parameters()),
lr=args.lr)
loss_class = nn.CrossEntropyLoss().cuda(cuda)
loss_val = None
# Check directories
check_dir_exists([
'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir
])
total, correct, top5correct, cnt = 0, 0, 0, 0
print('Start training ...')
for epoch in range(args.epoch):
# Evaluation cover
if epoch % 5 == 0 and epoch != 0:
mol.eval()
eval_dir = os.path.join(evaluation_dir, 'epoch%d' % epoch)
evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda,
eval_dir, args)
encode_accuracy, encode_top5accuracy, fc_accuracy, fc_top5accuracy = evaluate_labeled_data(
test_loader, mol, cuda)
writer.add_scalar('test/encode_feature_accuracy',
np.mean(encode_accuracy), epoch)
writer.add_scalar('test/encode_feature_top5accuracy',
np.mean(encode_top5accuracy), epoch)
writer.add_scalar('test/fc_feature_accuracy', np.mean(fc_accuracy),
epoch)
writer.add_scalar('test/fc_feature_top5accuracy',
np.mean(fc_top5accuracy), epoch)
if epoch != 0:
# Testing classifier
mol.eval()
test_acc, test_top5acc = test(test_loader, mol, cuda, 'Full')
writer.add_scalar('test/class_accuracy', test_acc, epoch)
writer.add_scalar('test/class_top5accuracy', test_top5acc, epoch)
step_time = time.time()
mol.train()
print('######### Training with %d batches total ##########' %
len(train_loader))
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
label = Variable(
torch.Tensor([y[2][i] for i in range(len(y[0]))]).long())
label = label.cuda() if cuda else label
prob_class = mol(b_x)
loss = loss_class(prob_class, label)
writer.add_scalar('train/loss_classifier', loss, cnt)
optimizer1.zero_grad()
loss.backward()
optimizer1.step()
_, predicted = torch.max(prob_class.data, 1)
total += label.size(0)
correct += (predicted == label).sum().item()
top5pre = prob_class.topk(5, 1, True, True)
top5pre = top5pre[1].t()
top5correct += top5pre.eq(label.view(
1, -1).expand_as(top5pre)).sum().item()
writer.add_scalar('train/accuracy', correct / total, cnt)
writer.add_scalar('train/top5_accuracy', top5correct / total, cnt)
loss_val = 0.99 * loss_val + 0.01 * loss.data[
0] if loss_val is not None else loss.data[0]
if step % 10 == 0:
if os.path.exists(model_name):
shutil.copy2(model_name,
model_name.split('.pkl')[0] + '_back.pkl')
torch.save(mol, model_name)
print(
'[Training] Epoch:', epoch, 'Step:', step, '|',
'train loss %.6f; Time cost %.2f s; Accuracy %.3f%%; Top5 Accuracy %.3f%%'
% (loss.data[0], time.time() - step_time,
correct * 100 / total, top5correct * 100 / total))
correct, total, top5correct = 0, 0, 0
step_time = time.time()
cnt += 1
print('Finished. Totally cost %.2f' % (time.time() - start_time))
writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json'))
writer.close()
def train():
################################################################
# Arguments
################################################################
ae_args = train_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
start_time = time.time()
args = ae_args
model_name = 'model/vgg_classifier_%s%s_model-%s.pkl' % (
args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
evaluation_dir = 'res/evaluation_pic/vgg_%s%s-%s' % (
args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
vgg = torch.load(model_name).to(device)
else:
print('Init model ...')
vgg = VGGNet(args.fea_c).to(device)
# train_loader = getDataLoader(args, kwargs)
test_loader = getDataLoader(args, kwargs, train='test')
# cover_loader = getDataLoader(args, kwargs, train='cover')
cover_val_loader = getDataLoader(args, kwargs, train='cover_validation')
cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample')
optimizer = torch.optim.Adam(list(vgg.parameters()), lr=args.lr)
loss_class = nn.CrossEntropyLoss().cuda(cuda)
check_dir_exists(['res/', 'model', 'res/evaluation_pic', evaluation_dir])
# # Evaluation
# vgg.eval()
# check_dir_exists([os.path.join(evaluation_dir, 'cos'), os.path.join(evaluation_dir, 'distance')])
# # evaluate_cover(cover_val_loader, cover_sample_loader, vgg, cuda, evaluation_dir, args)
# fc_accuracy, fc_top5accuracy = evaluate_labeled_data(test_loader, vgg, cuda, both=False)
# # print('Encode accuracy:', np.mean(encode_accuracy))
# # print('Encode top5 accuracy:', np.mean(encode_top5accuracy))
# print('Fc accuracy:', np.mean(fc_accuracy))
# print('Fc top5 accuracy:', np.mean(fc_top5accuracy))
total, correct, top5correct, loss_total = 0, 0, 0, 0
t_per_img = []
for epoch in range(1):
step_time = time.time()
for step, (x, y) in enumerate(test_loader):
b_x = Variable(x, volatile=True).cuda() if cuda else Variable(x)
label = Variable(
torch.Tensor([y[2][i] for i in range(len(y[0]))]).long())
label = label.cuda() if cuda else label
t0 = time.time()
prob_class = vgg(b_x)
t_tmp = (time.time() - t0) / len(b_x) * 1000
t_per_img.append(t_tmp)
print(
'cost %.6fms per image this batch. cost %.6fms per image till now.'
% (t_tmp, np.mean(sorted(t_per_img)[1:-1])))
loss = loss_class(prob_class, label) # mean square error
# optimizer.zero_grad() # clear gradients for this training step
# loss.backward()
# optimizer.step()
_, predicted = torch.max(prob_class.data, 1)
total += label.size(0)
correct += (predicted == label).sum().item()
top5pre = prob_class.topk(5, 1, True, True)
top5pre = top5pre[1].t()
top5correct += top5pre.eq(label.view(
1, -1).expand_as(top5pre)).sum().item()
loss_total += loss.data[0] * label.size(0)
if step % 10 == 0:
torch.save(vgg, model_name)
print(
'Epoch:', epoch, 'Step:', step, '|',
'test loss %.6f; Time cost %.2f s; Classification error %.6f; '
'Top1 Accuracy %.3f; Top5 Accuracy %.3f' %
(loss_total / total, time.time() - step_time, loss.item(),
correct * 100 / total, top5correct * 100 / total))
step_time = time.time()
loss = None
print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train():
################################################################
# Arguments
################################################################
ae_args = args = feature_classifier_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
################################################################
# Load or Init model
################################################################
start_time = time.time()
model_name = 'model/FeaClass_%s_%s%s_model-%s.pkl' \
% (args.main_model, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
feature_name = 'model/%s_%s%s_model-%s.pkl' % (
args.main_model, args.model,
'' if args.fea_c is None else args.fea_c, args.dataset)
print('Loading feature model (%s)...' % feature_name)
assert os.path.exists(feature_name)
if cuda:
fea_mol = torch.load(feature_name).to(device)
else:
fea_mol = torch.load(feature_name, map_location='cpu').to(device)
mol = FClassifier(fea_mol, args.main_model, args.fea_c).to(device)
train_loader = getDataLoader(args, kwargs)
optimizer = torch.optim.Adam(list(mol.classification.parameters()),
lr=args.lr)
loss_fn = nn.CrossEntropyLoss().cuda(cuda)
check_dir_exists(['res/', 'model'])
total, correct, top5correct = 0, 0, 0
for epoch in range(args.epoch):
step_time = time.time()
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
label = Variable(
torch.Tensor([y[2][i] for i in range(len(y[0]))]).long())
label = label.cuda() if cuda else label
prob_class = mol(b_x)
loss = loss_fn(prob_class, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
_, predicted = torch.max(prob_class.data, 1)
total += label.size(0)
correct += (predicted == label).sum().item()
top5pre = prob_class.topk(5, 1, True, True)
top5pre = top5pre[1].t()
top5correct += top5pre.eq(label.view(
1, -1).expand_as(top5pre)).sum().item()
if step % 10 == 0:
shutil.copy2(model_name,
model_name.split('.pkl')[0] + '_back.pkl')
torch.save(mol, model_name)
print(
'Epoch:', epoch, 'Step:', step, '|',
'train loss %.6f; Time cost %.2f s; Classification error %.6f'
'Accuracy %.3f%%; Top5 Accuracy %.3f%%' %
(loss.data[0], time.time() - step_time, loss,
correct * 100 / total, top5correct * 100 / total))
correct, total, top5correct = 0, 0, 0
step_time = time.time()
print('Finished. Totally cost %.2f' % (time.time() - start_time))
def train(args, mol_short='AEClass_both', main_model=AEClass):
print('######### This is train function #############')
################################################################
# Arguments
################################################################
ae_args = args
if args.gpu == -1:
os.environ["CUDA_VISIBLE_DEVICES"] = GPU().choose_gpu()
elif args.gpu is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
d_name = '' if args.decoder == 'vgg' else '_' + args.decoder + 'decoder'
mol_short = mol_short + args.name + d_name
log_dir = 'log/log_%s_%s%s_model-%s/' %\
(mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
remove_dir_exists([log_dir])
writer = SummaryWriter(log_dir)
start_time = time.time()
model_name = 'model/%s_%s%s_model-%s.pkl' % (mol_short, args.model,
'' if args.fea_c is None else
args.fea_c, args.dataset)
print('[Model] model name is', model_name)
pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model,
'' if args.fea_c is None else args.fea_c,
args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c,
args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
if cuda:
mol = torch.load(model_name).to(device)
else:
mol = torch.load(model_name, map_location='cpu')
else:
print('Init model ...')
mol = main_model(args.fea_c).to(device)
print('Prepare data loader ...')
# train_loader = getDataLoader(args, kwargs, train='train')
fuse_loader = getDataLoader(args, kwargs, train='fuse', p=args.imgnet_p)
test_dataset = RandomBatchSampler(getDataset(args, train='test'),
args.batch_size)
test_loader = getDataLoader(args, kwargs, train='test')
# small_test_loader = getDataLoader(args, kwargs, train=False, p=10)
# cover_loader = getDataLoader(args, kwargs, train='cover')
cover_val_loader = getDataLoader(args, kwargs, train='cover_validation')
cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample')
# optimizer_cls = torch.optim.Adam(get_optimized_params(mol, 'classifier', lr=args.lr) +
# get_optimized_params(mol, 'small_features', lr=args.lr) +
# get_optimized_params(mol, 'vgg.38', lr=args.lr/5), lr=args.lr)
optimizer_d = torch.optim.Adam(
([] if args.fea_c == 512 else get_optimized_params(
mol, 'small_features', args.lr)) +
get_optimized_params(mol, 'decoder', args.lr) +
get_optimized_params(mol, 'vgg.38', args.lr / 5),
lr=args.lr)
optimizer = torch.optim.Adam(
([] if args.fea_c == 512 else get_optimized_params(
mol, 'small_features', args.lr)) +
get_optimized_params(mol, 'classifier', lr=args.lr) +
get_optimized_params(mol, 'decoder', args.lr) +
get_optimized_params(mol, 'vgg.38', args.lr / 5),
lr=args.lr)
loss_decoder = nn.MSELoss()
loss_class = nn.CrossEntropyLoss(reduction='none').cuda(cuda)
def loss_cls_fn(prob_class, label, weight, beta=0.01):
if torch.sum(weights) == 0:
return 99999
else:
lss = loss_class(prob_class, label) * weight
lss = torch.sum(lss) / torch.sum(weight)
return lss * beta
check_dir_exists([
'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir
])
total, correct, top5correct, cnt = 0, 0, 0, 0
print('Start training ...')
for epoch in range(args.epoch):
if (epoch % 5 == 0) and epoch != 0:
# Evaluation on cover data
mol.eval()
eval_dir = os.path.join(evaluation_dir, 'epoch%d' % epoch)
evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda,
eval_dir, args)
fc_accuracy, fc_top5accuracy = evaluate_labeled_data(test_loader,
mol,
cuda,
both=False)
print('Fc accuracy:', np.mean(fc_accuracy))
print('Fc top5 accuracy:', np.mean(fc_top5accuracy))
if epoch >= 0:
mol.eval()
# Testing on ImageNet val
print('######### Testing on ImageNet val Dataset ###########')
test_loss_decoder, test_loss_cls, test_acc, test_top5acc = test_cls_decoder(
test_loader, mol, cuda, 'Full')
# test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001
writer.add_scalar('test_imagenet/loss_decoder', test_loss_decoder,
epoch)
writer.add_scalar('test_imagenet/loss_classifier', test_loss_cls,
epoch)
# writer.add_scalar('test_imagenet/loss', test_loss, epoch)
writer.add_scalar('test_imagenet/accuracy', test_acc, epoch)
writer.add_scalar('test_imagenet/top5accuracy', test_top5acc,
epoch)
# Testing on Cover val
print('######### Testing on Cover val Dataset ###########')
# test_loss_decoder, test_loss_cls, test_acc, test_top5acc = test_cls_decoder(cover_val_loader, mol, cuda, 'Full')
test_loss_decoder = test_decoder(cover_val_loader, mol, cuda,
'Full')
# test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001
writer.add_scalar('test_cover/loss_decoder', test_loss_decoder,
epoch)
# writer.add_scalar('test_cover/loss_classifier', test_loss_cls, epoch)
# writer.add_scalar('test_cover/loss', test_loss, epoch)
# writer.add_scalar('test_cover/accuracy', test_acc, epoch)
# writer.add_scalar('test_cover/top5accuracy', test_top5acc, epoch)
step_time = time.time()
for step, (x, y) in enumerate(fuse_loader):
mol.train()
b_x = Variable(x).cuda() if cuda else Variable(x)
b_y = b_x.detach().cuda() if cuda else b_x.detach(
) # batch y, shape (batch, 32*32*3)
# //TODO
# train_cls =
label = Variable(
torch.Tensor([y[2][i] for i in range(len(y[0]))]).long())
label = label.cuda() if cuda else label
weights = Variable(
torch.Tensor([y[3][i] for i in range(len(y[0]))]))
weights = weights.cuda() if cuda else weights
_, decoded, prob_class = mol(b_x)
if step % 100 == 0:
img_to_save = decoded.data
save_image(img_to_save,
'%s/%s-%s.jpg' % (pic_dir, epoch, step))
loss1 = loss_decoder(decoded, b_y)
loss2 = loss_cls_fn(prob_class, label, weights, beta=0.01)
loss = 0
writer.add_scalar('train/loss_decoder', loss1, cnt)
writer.add_scalar('train/loss_classifier', loss2, cnt)
# writer.add_scalar('train/loss', loss, cnt)
if loss2 != 99999:
loss = (1 - args.alpha) * loss2 + args.alpha * loss1
optimizer.zero_grad()
loss.backward()
optimizer.step()
_, predicted = torch.max(prob_class.data, 1)
predicted = torch.Tensor([
predicted[i] for i in range(len(weights))
if weights[i] == 1
])
label2 = torch.Tensor(
[label[i] for i in range(len(weights)) if weights[i] == 1])
total += predicted.shape[0]
correct += (predicted == label2).sum().item()
prob_class2 = prob_class[[
i for i in range(len(weights)) if weights[i] == 1
]]
top5pre = prob_class2.topk(5, 1, True, True)
top5pre = top5pre[1].t()
label2 = label2.long().cuda() if cuda else label2.long()
top5correct += top5pre.eq(
label2.view(1, -1).expand_as(top5pre)).sum().item()
writer.add_scalar('train/accuracy', correct / total, cnt)
writer.add_scalar('train/top5_accuracy', top5correct / total,
cnt)
else:
optimizer_d.zero_grad()
loss1.backward()
optimizer_d.step()
# loss_val = 0.99*loss_val + 0.01*loss.data[0] if loss_val is not None else loss.data[0]
if step % 10 == 0:
if os.path.exists(model_name):
shutil.copy2(model_name,
model_name.split('.pkl')[0] + '_back.pkl')
torch.save(mol, model_name)
total_tmp = total if total != 0 else 1
t_x, t_y = test_dataset.get_sample()
test_acc, test_top5_acc = _test_sample_batch(
mol, cuda, t_x, t_y)
print(
'[Training] Epoch:', epoch, 'Step:', step, '|',
'Time cost %.2f s; Classification error %.6f; Decoder error %.6f; Loss %.6f; '
'Accuracy %.3f%%/%.3f%%; Top5 Accuracy %.3f%%/%.3f%% (%s)'
% (time.time() - step_time, 0.0888
if type(loss2) == int else loss2.data[0], loss1.data[0],
0.0888 if type(loss) == int else loss.data[0],
correct * 100 / total_tmp, test_acc * 100, top5correct *
100 / total_tmp, test_top5_acc * 100, time.ctime()))
writer.add_scalar('test/accuracy', test_acc, cnt)
writer.add_scalar('test/top5_accuracy', test_top5_acc, cnt)
correct, total, top5correct = 0, 0, 0
step_time = time.time()
cnt += 1
print('Finished. Totally cost %.2f' % (time.time() - start_time))
writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json'))
writer.close()
def train_decoder_only(args, mol_short='AEClass_d', main_model=AEClass):
print('######### This is train decoder only function #############')
################################################################
# Arguments
################################################################
ae_args = args
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
d_name = '' if args.decoder == 'vgg' else '_' + args.decoder + 'decoder'
mol_short = mol_short + args.name + d_name
log_dir = 'log/log_%s_%s%s_model-%s/' %\
(mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
remove_dir_exists([log_dir])
writer = SummaryWriter(log_dir)
start_time = time.time()
model_name = 'model/%s_%s%s_model-%s.pkl' \
% (mol_short, args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
print('[Model] model name is', model_name)
pic_dir = 'res/%s_%s%s-%s/' % (mol_short, args.model,
'' if args.fea_c is None else args.fea_c,
args.dataset)
evaluation_dir = 'res/evaluation_pic/%s_%s%s-%s' % (
mol_short, args.model, '' if args.fea_c is None else args.fea_c,
args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
mol = main_model(args.fea_c, decoder=args.decoder).to(device)
print('Prepare data loader ...')
test_loader = getDataLoader(args, kwargs, train='test')
# small_test_loader = getDataLoader(args, kwargs, train=False, p=10)
cover_loader = getDataLoader(args, kwargs, train='cover')
cover_val_loader = getDataLoader(args, kwargs, train='cover_validation')
cover_sample_loader = getDataLoader(args, kwargs, train='cover_sample')
# optimizer1 = torch.optim.Adam(list(mol.classification.parameters())+list(mol.small_features.parameters())+
# list(mol.decoder.parameters()), lr=args.lr)
optimizer1 = torch.optim.Adam(
([] if args.fea_c == 512 else get_optimized_params(
mol, 'small_features', args.lr)) +
get_optimized_params(mol, 'decoder', args.lr) +
get_optimized_params(mol, 'vgg.38', args.lr / 5),
lr=args.lr)
loss_decoder = nn.MSELoss()
check_dir_exists([
'res/', 'model', pic_dir, log_dir, 'res/evaluation_pic', evaluation_dir
])
loss_val = None
print('Start training ...')
cnt = 0
for epoch in range(args.epoch):
if epoch % 5 == (5 - 1) or epoch == (args.epoch - 1):
# Evaluation on cover data
eval_dir = os.path.join(evaluation_dir, 'epoch%d' % epoch)
evaluate_cover(cover_val_loader, cover_sample_loader, mol, cuda,
eval_dir)
if epoch != 0:
# Testing on Cover val
print('######### Testing on Cover val Dataset ###########')
test_loss_decoder = test_decoder(cover_val_loader, mol, cuda,
'Full')
# test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001
writer.add_scalar('test_cover/loss_decoder', test_loss_decoder,
epoch)
# writer.add_scalar('test_cover/loss_classifier', test_loss_cls, epoch)
# writer.add_scalar('test_cover/loss', test_loss, epoch)
# writer.add_scalar('test_cover/accuracy', test_acc, epoch)
# writer.add_scalar('test_cover/top5accuracy', test_top5acc, epoch)
# Testing on ImageNet val
print('######### Testing on ImageNet val Dataset ###########')
# test_loss_decoder, test_loss_cls, test_acc, test_top5acc = test_cls_decoder(test_loader, mol, cuda, 'Full')
test_loss_decoder = test_decoder(test_loader, mol, cuda, 'Full')
# test_loss = (1 - args.alpha) * test_loss_cls + args.alpha * test_loss_decoder / 0.001
writer.add_scalar('test_imagenet/loss_decoder', test_loss_decoder,
epoch)
# writer.add_scalar('test_imagenet/loss_classifier', test_loss_cls, epoch)
# writer.add_scalar('test_imagenet/loss', test_loss, epoch)
# writer.add_scalar('test_imagenet/accuracy', test_acc, epoch)
# writer.add_scalar('test_imagenet/top5accuracy', test_top5acc, epoch)
step_time = time.time()
print('######### Training with %d batches total ##########' %
len(cover_loader))
for step, (x, y) in enumerate(cover_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
b_y = b_x.detach().cuda() if cuda else b_x.detach(
) # batch y, shape (batch, 32*32*3)
_, decoded, prob_class = mol(b_x)
if step % 500 == 0 or step == (len(cover_loader) - 1):
img_to_save = decoded.data
save_image(img_to_save,
'%s/%s-%s.jpg' % (pic_dir, epoch, step))
loss1 = loss_decoder(decoded, b_y)
# loss = (1-args.alpha) * loss2 + args.alpha * loss1 / 0.001
loss = loss1
writer.add_scalar('train/loss_decoder', loss1, cnt)
writer.add_scalar('train/loss', loss, cnt)
optimizer1.zero_grad()
loss.backward()
optimizer1.step()
loss_val = 0.99 * loss_val + 0.01 * loss.data[
0] if loss_val is not None else loss.data[0]
if step % 10 == 0 or step == (len(cover_loader) - 1):
if os.path.exists(model_name):
shutil.copy2(model_name,
model_name.split('.pkl')[0] + '_back.pkl')
torch.save(mol, model_name)
print(
'[Training] Epoch:', epoch, 'Step:', step, '|',
'train loss %.6f; Time cost %.2f s; Decoder error %.6f;' %
(loss.data[0], time.time() - step_time, loss1.data[0]))
step_time = time.time()
cnt += 1
print('Finished. Totally cost %.2f' % (time.time() - start_time))
writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json'))
writer.close()
def train():
################################################################
# Arguments
################################################################
ae_args = train_args()
cuda = ae_args.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if ae_args.cuda else {}
# global ae_args, cuda, device, kwargs
args = ae_args
log_dir = 'log/log_AE_%s%s_model-%s/' %\
(args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
writer = SummaryWriter(log_dir)
start_time = time.time()
model_name = 'model/AE_%s%s_model-%s.pkl' % (
args.model, '' if args.fea_c is None else args.fea_c, args.dataset)
pic_dir = 'res/AE_%s%s-%s/' % (args.model, '' if args.fea_c is None else
args.fea_c, args.dataset)
if os.path.exists(model_name) and args.load_model:
print('Loading model ...')
mol = torch.load(model_name).to(device)
else:
print('Init model ...')
mol = AutoEncoder(args.fea_c).to(device)
print('Prepare data loader ...')
train_loader = getDataLoader(args, kwargs, train=True)
test_loader = getDataLoader(args, kwargs, train=False)
optimizer1 = torch.optim.Adam(list(mol.small_features.parameters()) +
list(mol.decoder.parameters()),
lr=args.lr)
optimizer2 = torch.optim.Adam(list(mol.features.parameters()),
lr=args.lr / 10)
loss_decoder = nn.MSELoss()
check_dir_exists(['res/', 'model', pic_dir, log_dir])
total, correct, top5correct, cnt = 0, 0, 0, 0
print('Start training ...')
for epoch in range(args.epoch):
# Testing
test_acc, test_top5acc, test_loss = test_feature(
test_loader, mol, cuda, 'Full')
writer.add_scalar('test/accuracy', np.mean(test_acc), epoch)
writer.add_scalar('test/top5accuracy', np.mean(test_top5acc), epoch)
writer.add_scalar('test/loss_decoder', test_loss, epoch)
step_time = time.time()
for step, (x, y) in enumerate(train_loader):
b_x = Variable(x).cuda() if cuda else Variable(x)
b_y = b_x.detach().cuda() if cuda else b_x.detach(
) # batch y, shape (batch, 32*32*3)
encoded, decoded = mol(b_x)
if step % 100 == 0:
img_to_save = decoded.data
save_image(img_to_save,
'%s/%s-%s.jpg' % (pic_dir, epoch, step))
loss = loss_decoder(decoded, b_y)
writer.add_scalar('train/loss_decoder', loss, cnt)
optimizer1.zero_grad()
optimizer2.zero_grad()
loss.backward()
optimizer1.step()
optimizer2.step()
if step % 50 == 0:
if os.path.exists(model_name):
shutil.copy2(model_name,
model_name.split('.pkl')[0] + '_back.pkl')
torch.save(mol, model_name)
print(
'[Training] Epoch:', epoch, 'Step:', step, '|',
'train loss %.6f; Time cost %.2f s; Decoder error %.6f' %
(loss.data[0], time.time() - step_time, loss))
step_time = time.time()
cnt += 1
print('Finished. Totally cost %.2f' % (time.time() - start_time))
writer.export_scalars_to_json(os.path.join(log_dir, 'all_scalars.json'))
writer.close()
