def train_epoch_wo_outlier(model, optimizer, in_loader, loss_func, cur_epoch, op_cfg, writer):
global global_cfg
model.train()
avg_loss = 0
correct = 0
in_data_size = len(in_loader.dataset)
for cur_iter, in_set in enumerate(in_loader):
#TODO: Dimension of in_set and out_set should be checked!
# Data to GPU
data = in_set[0]
targets = in_set[1]
if cur_iter == 0:
writer.add_image('in_dist target {}'.format(targets[0]), data[0], cur_epoch)
data, targets = data.cuda(), targets.cuda()
# Adjust Learning rate
lr = optim.get_lr_at_epoch(op_cfg, cur_epoch + float(cur_iter) / in_data_size)
optim.set_lr(optimizer, lr)
# Foward propagation and Calculate loss
logits = model(data)
(ava_logits, ova_logits) = logits
#print(logits.size())
global_cfg['loss']['model'] = model
global_cfg['loss']['data'] = data
loss_dict = loss_func(logits, targets, global_cfg['loss'])
loss = loss_dict['loss']
logits = F.softmax(ava_logits, dim=1)
# Back propagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Calculate classifier error about in-distribution sample
num_topks_correct = metrics.topks_correct(logits[:len(targets)], targets, (1,))
[top1_correct] = [x for x in num_topks_correct]
# Add additional metrics!!!
loss, top1_correct = loss.item(), top1_correct.item()
avg_loss += loss
correct += top1_correct
summary = {
'avg_loss': avg_loss / in_data_size,
'classifier_acc': correct / in_data_size,
'lr': optim.get_lr_at_epoch(op_cfg, cur_epoch),
'epoch': cur_epoch,
}
return summary
def train_epoch_wo_outlier(model, optimizer, in_loader, attack_in, cur_epoch, op_cfg, writer):
global global_cfg
model.train()
avg_loss = 0
correct = 0
in_data_size = len(in_loader.dataset)
for cur_iter, (x_tf_0, x_tf_90, x_tf_180, x_tf_270, targets) in enumerate(in_loader):
batch_size = x_tf_0.shape[0]
assert x_tf_0.shape[0] == \
x_tf_90.shape[0] == \
x_tf_180.shape[0] == \
x_tf_270.shape[0] == \
targets.shape[0]
#x_tf_trans.shape[0] == \
#target_trans_x.shape[0] == \
#target_trans_y.shape[0] == \
batch = np.concatenate((
x_tf_0,
x_tf_90,
x_tf_180,
x_tf_270
), 0)
batch = torch.FloatTensor(batch).cuda()
target_rots = torch.cat((
torch.zeros(batch_size),
torch.ones(batch_size),
2 * torch.ones(batch_size),
3 * torch.ones(batch_size)
), 0).long()
if attack_in is not None:
# Process PGD attack
batch = attack_in.perturb(batch, batch_size, torch.cat((targets, target_rots), 0).cuda())
batch = batch.cuda()
if cur_iter == 0:
writer.add_image('Original', batch[0], cur_epoch)
writer.add_image('Rot90', batch[batch_size], cur_epoch)
writer.add_image('Rot180', batch[batch_size * 2], cur_epoch)
writer.add_image('Rot270', batch[batch_size * 3], cur_epoch)
# Adjust Learning rate
lr = optim.get_lr_at_epoch(op_cfg, cur_epoch + float(cur_iter) / in_data_size)
optim.set_lr(optimizer, lr)
logits, pen = model(batch)
classification_logits = logits[:batch_size]
rot_logits = model.rot_head(pen[:4*batch_size])
#x_trans_logits = model.x_trans_head(pen[4*batch_size:])
#y_trans_logits = model.y_trans_head(pen[4*batch_size:])
classification_loss = F.cross_entropy(classification_logits, targets.cuda())
rot_loss = F.cross_entropy(rot_logits, target_rots.cuda()) * global_cfg['loss']['rot_weight']
# x_trans_loss = F.cross_entropy(x_trans_logits, target_trans_x.cuda()) * global_cfg['loss']['trans_weight']
# y_trans_loss = F.cross_entropy(y_trans_logits, target_trans_y.cuda()) * global_cfg['loss']['trans_weight']
#loss = classification_loss + ((rot_loss + x_trans_loss + y_trans_loss) / 3.0)
loss = classification_loss + rot_loss
# Back propagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Calculate classifier error about in-distribution sample
num_topks_correct = metrics.topks_correct(logits[:batch_size], targets.cuda(), (1,))
[top1_correct] = [x for x in num_topks_correct]
# Add additional metrics!!!
loss, top1_correct = loss.item(), top1_correct.item()
avg_loss += loss
correct += top1_correct
summary = {
'avg_loss': avg_loss / in_data_size,
'classifier_acc': correct / in_data_size,
'lr': optim.get_lr_at_epoch(op_cfg, cur_epoch),
'epoch': cur_epoch,
}
return summary
def train_epoch_w_outlier(model, optimizer, in_loader, out_loader, loss_func, detector_func, cur_epoch, op_cfg, writer):
global global_cfg
model.train()
avg_loss = 0
correct = 0
total = 0
in_data_size = len(in_loader.dataset)
out_loader.dataset.offset = np.random.randint(len(out_loader.dataset))
for cur_iter, (in_set, out_set) in enumerate(zip(in_loader, out_loader)):
#TODO: Dimension of in_set and out_set should be checked!
# Data to GPU
data = torch.cat((in_set[0], out_set[0]), 0)
targets = in_set[1]
if cur_iter == 0:
writer.add_image('in_dist sample, target:[{}]'.format(targets[0]), in_set[0][0], cur_epoch)
writer.add_image('out_dist sample', out_set[0][0], cur_epoch)
data, targets = data.cuda(), targets.cuda()
# Adjust Learning rate
lr = optim.get_lr_at_epoch(op_cfg, cur_epoch + float(cur_iter) / in_data_size)
optim.set_lr(optimizer, lr)
# Foward propagation and Calculate loss and confidence
logits = model(data)
global_cfg['loss']['model'] = model
global_cfg['loss']['data'] = data
global_cfg['detector']['model'] = model
global_cfg['detector']['data'] = data
loss_dict = loss_func(logits, targets, global_cfg['loss'])
loss = loss_dict['loss']
confidences_dict = detector_func(logits, targets, global_cfg['detector'])
confidences = confidences_dict['confidences']
# Back propagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
## METRICS ##
# Calculate classifier error about in-distribution sample
num_topks_correct = metrics.topks_correct(logits[:len(targets)], targets, (1,))
[top1_correct] = [x for x in num_topks_correct]
# Calculate OOD metrics (auroc, aupr, fpr)
#(auroc, aupr, fpr) = metrics.get_ood_measures(confidences, targets)
# Add additional metrics!!!
## UDATE STATS ##
loss, top1_correct = loss.item(), top1_correct.item()
avg_loss += loss
correct += top1_correct
total += targets.size(0)
summary = {
'avg_loss': avg_loss / total,
'classifier_acc': correct / total,
'lr': optim.get_lr_at_epoch(op_cfg, cur_epoch),
'epoch': cur_epoch,
}
return summary
def train_epoch_wo_outlier(model, optimizer, in_loader, loss_func, cur_epoch,
op_cfg, writer):
global global_cfg
model.train()
avg_loss = 0
avg_sup_loss = 0
avg_con_loss = 0
correct = 0
in_data_size = len(in_loader.dataset)
for cur_iter, in_set in enumerate(in_loader):
#TODO: Dimension of in_set and out_set should be checked!
# Data to GPU
data = in_set[0]
if cur_iter == 0:
writer.add_image('Original', data[0], cur_epoch)
# Transform imgs
data = np.transpose((data.numpy() * 255).round().astype(np.uint8),
(0, 2, 3, 1))
images_aug = seq0(images=data)
data0 = torch.from_numpy(
np.transpose((np.stack(images_aug, 0).astype(np.float32) / 255.),
(0, 3, 1, 2)))
images_aug = seq1(images=data)
data1 = torch.from_numpy(
np.transpose((np.stack(images_aug, 0).astype(np.float32) / 255.),
(0, 3, 1, 2)))
if cur_iter == 0:
writer.add_image('Transform0', data0[0], cur_epoch)
writer.add_image('Transform1', data1[0], cur_epoch)
data = torch.cat((data0, data1), 0)
targets = in_set[1]
data, targets = data.cuda(), targets.cuda()
# Adjust Learning rate
lr = optim.get_lr_at_epoch(op_cfg,
cur_epoch + float(cur_iter) / in_data_size)
optim.set_lr(optimizer, lr)
# Foward propagation and Calculate loss
(g_logits, h_logits, f_logits) = model(data, cur_epoch)
# logits: (g, h)
logits = (g_logits, h_logits)
global_cfg['loss']['model'] = model
global_cfg['loss']['data'] = data
loss_dict = loss_func(logits, targets, global_cfg['loss'])
loss = loss_dict['loss']
sup_loss = loss_dict['sup_loss']
con_loss = loss_dict['con_loss']
# Back propagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Calculate classifier error about in-distribution sample
num_topks_correct = metrics.topks_correct(g_logits[:len(targets)],
targets, (1, ))
[top1_correct] = [x for x in num_topks_correct]
# Add additional metrics!!!
loss, top1_correct = loss.item(), top1_correct.item()
avg_loss += loss
avg_sup_loss += sup_loss
avg_con_loss += con_loss
correct += top1_correct
summary = {
'avg_loss': avg_loss / in_data_size,
'avg_sup_loss': avg_sup_loss / in_data_size,
'avg_con_loss': avg_con_loss / in_data_size,
'classifier_acc': correct / in_data_size,
'lr': optim.get_lr_at_epoch(op_cfg, cur_epoch),
'epoch': cur_epoch,
}
return summary
def train_epoch_wo_outlier(feature_extractor, G, D, G_optimizer, D_optimizer,
in_loader, cur_epoch, op_cfg, writer):
global global_cfg
D.train()
G.train()
feature_extractor.eval()
avg_loss = 0
correct = 0
in_data_size = len(in_loader.dataset)
real_feature = None
real_target = None
for cur_iter, in_set in enumerate(in_loader):
#TODO: Dimension of in_set and out_set should be checked!
# Data to GPU
real_data = in_set[0]
targets = in_set[1]
if cur_iter == 0:
writer.add_image('in_dist target {}'.format(targets[0]),
real_data[0], cur_epoch)
real_data, targets = real_data.cuda(), targets.cuda()
_, real_features = feature_extractor(real_data, -1)
if cur_iter == 0:
real_feature = real_features[0].unsqueeze(0)
real_target = targets[0]
# Adjust Learning rate
lr = optim.get_lr_at_epoch(op_cfg,
cur_epoch + float(cur_iter) / in_data_size)
optim.set_lr(G_optimizer, lr)
optim.set_lr(D_optimizer, lr)
###
d_out_real, dr1, dr2 = D(real_data, real_features)
if global_cfg['loss']['adv_loss'] == 'wgan-gp':
d_loss_real = -torch.mean(d_out_real)
elif global_cfg['loss']['adv_loss'] == 'hinge':
d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()
z = tensor2var(torch.randn(real_data.size(0), G.z_dim))
fake_images, gf1, gf2 = G(z, real_features)
d_out_fake, df1, df2 = D(fake_images, real_features)
if global_cfg['loss']['adv_loss'] == 'wgan-gp':
d_loss_fake = d_out_fake.mean()
elif global_cfg['loss']['adv_loss'] == 'hinge':
d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()
# Backward + Optimize
d_loss = d_loss_real + d_loss_fake
if global_cfg['loss']['adv_loss'] == 'wgan-gp':
d_loss += op_cfg['lambda_gp'] * compute_gradient_penalty(
feature_extractor, D, real_data.data, fake_images.data)
D_optimizer.zero_grad()
G_optimizer.zero_grad()
d_loss.backward()
D_optimizer.step()
# ================== Train G and gumbel ================== #
# Create random noise
z = tensor2var(torch.randn(real_data.size(0), G.z_dim))
_, real_features = feature_extractor(real_data, -1)
fake_images, _, _ = G(z, real_features)
# Compute loss with fake images
g_out_fake, _, _ = D(fake_images, real_features) # batch x n
_, fake_features = feature_extractor(fake_images, -1)
if global_cfg['loss']['adv_loss'] == 'wgan-gp':
g_loss_fake = -g_out_fake.mean()
elif global_cfg['loss']['adv_loss'] == 'hinge':
g_loss_fake = -g_out_fake.mean()
g_loss_feature = F.mse_loss(fake_features, real_features)
g_loss = g_loss_fake + g_loss_feature
D_optimizer.zero_grad()
G_optimizer.zero_grad()
g_loss.backward()
G_optimizer.step()
###
# Add additional metrics!!!
avg_loss += (d_loss + g_loss)
## Epoch
# Print out log info
summary = {
'avg_loss': avg_loss / in_data_size,
'lr': optim.get_lr_at_epoch(op_cfg, cur_epoch),
'epoch': cur_epoch,
'real_feature': real_feature,
'real_target': real_target,
}
return summary
def main(cfg):
# basic settings
loss_F = torch.nn.CrossEntropyLoss()
gpu_nums = int(cfg['NUM_GPUS'])
if gpu_nums == 0:
use_cuda = False
else:
use_cuda = True
# load model
model = AnyNet(cfg)
if use_cuda:
model = torch.nn.DataParallel(model, device_ids=[0])
model = model.cuda()
# load_dataset
Trainpath = cfg['TRAIN']['PATH']
RESIZE_SIZE = cfg['TRAIN']['IM_SIZE']
train_data = SingleDataset(Trainpath, split='train', resize_size=RESIZE_SIZE)
train_loader= DataLoader(dataset=train_data, batch_size=cfg['TRAIN']['BATCH_SIZE'],
shuffle=True, num_workers=cfg['DATA_LOADER']['NUM_WORKERS'], pin_memory=True)
Testpath = cfg['TEST']['PATH']
RESIZE_SIZE_val = cfg['TEST']['IM_SIZE']
test_data = SingleDataset(Testpath, split='val', resize_size=RESIZE_SIZE_val)
test_loader = DataLoader(dataset=test_data, batch_size=cfg['TEST']['BATCH_SIZE'],
shuffle=False, num_workers=cfg['DATA_LOADER']['NUM_WORKERS'], pin_memory=True)
# optimizer and loss function and evaluator
if cfg['OPTIM']['OPTIMIZER'] == 'Adam':
optimizer = torch.optim.Adam(model.parameters(), lr=cfg['OPTIM']['BASE_LR'], weight_decay=1e-4)
else:
optimizer = torch.optim.SGD(model.parameters(), lr=cfg['OPTIM']['BASE_LR'], momentum=0.9, weight_decay=5e-4)
# load checkpoint or initial weights
start_epoch = 0
if cfg['TRAIN']['RESUME'] is not None:
resume = cfg['TRAIN']['RESUME']
if not os.path.isfile(resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(resume))
checkpoint_epoch = cp.load_checkpoint(resume, gpu_num=gpu_nums, model=model, optimizer=optimizer)
start_epoch = checkpoint_epoch + 1
elif cfg['TRAIN']['WEIGHTS']:
cp.load_checkpoint(cfg['TRAIN']['WEIGHTS'], gpu_nums, model)
else:
init_weights(model, zero_init_gamma=cfg['BN']['ZERO_INIT_FINAL_GAMMA'])
# save training process
log_file = log_g.get_log_filename(os.path.join(cfg['OUT_DIR'], 'log/'))
log = open(log_file, 'w+')
# start training
max_epoch = cfg['OPTIM']['MAX_EPOCH']
batch_size = cfg['TRAIN']['BATCH_SIZE']
eval_period = cfg['TRAIN']['EVAL_PERIOD']
batch_count = 0
total_step = len(train_loader)
num_class = cfg['MODEL']['NUM_CLASSES']
# correct_all = list(0. for i in range(cfg['MODEL']['NUM_CLASSES']))
# total_all = list(0. for i in range(cfg['MODEL']['NUM_CLASSES']))
for epoch in range(start_epoch, max_epoch):
print('**************train --%d-- **************' % (epoch))
log.write('**************train --%d-- **************\n' % (epoch))
# update learning rate
lr = optim.get_epoch_lr(epoch_i=epoch, cfg=cfg)
optim.set_lr(optimizer, lr)
#############################################################################
# start training an epoch
#############################################################################
model.train()
c_train = 0
t_train = 0
for i, (img, lbl) in enumerate(train_loader):
batch_count += 1
# use cuda
if use_cuda:
img, lbl = img.cuda(), lbl.cuda()
# forward
preds = model(img)
loss = loss_F(preds, lbl)
# backward
# optimizer.zero_grad()
loss.backward()
# optimizer.step()
torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
if (batch_count % batch_size) == 0:
optimizer.step()
optimizer.zero_grad()
batch_count = 0
_, predicted = preds.max(1)
c_train += predicted.eq(lbl).sum().item()
t_train += lbl.size(0)
# print epoch, step, loss, lr
print('[%s]--train: %d/%d\tstep:%d/%d----lr:%.5f---loss:%.4f---Acc:%.3f' % (
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
(epoch + 1), max_epoch, (i + 1), total_step, lr, loss.item(), 100*(c_train/t_train)))
log.write('[%s]--train: [%d/%d]\tstep: [%d/%d]\t----lr:%.5f---loss:%.4f---Acc:%.3f\n' %(
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
(epoch + 1), max_epoch, (i + 1), total_step, lr, loss.item(), 100*(c_train/t_train)))
#############################################################################
# start validation
#############################################################################
if ((epoch+1) % eval_period == 0):
print('**************validation --%d-- **************' % ((epoch + 1) // eval_period))
model.eval()
mean_loss_val = 0
correct = np.zeros((num_class))
total = np.zeros((num_class))
top1_acc_sum = []
with torch.no_grad():
for val_epoch, (img_val, lbl_val) in enumerate(test_loader):
if use_cuda:
img_val, lbl_val = img_val.cuda(), lbl_val.cuda()
# predict
preds_val = model(img_val)
# calculate loss
loss_val = loss_F(preds_val, lbl_val)
mean_loss_val += loss_val.item()
# evaluation
top1_acc, top2_acc = Evaluator.accuracy(preds_val, lbl_val, [1,2])
correct_i, total_i = Evaluator.accuracy_perclass(preds_val, lbl_val, num_class)
correct += correct_i
total += total_i
top1_acc_sum.append(top1_acc)
print('[%s]--valid: [%d/%d]\tloss: %.4f---top1_acc: %.3f' % (
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
val_epoch, len(test_loader), loss_val.item(), top1_acc.item()))
print('[{}]--valid: [{}]\tmean_loss: {}\ttop1_acc: {}\tper_class_acc: {}'.format(
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
(epoch + 1), (mean_loss_val / len(test_loader)), np.mean(top1_acc_sum), 100*(correct/total)))
# save log
log.write('[{}]--valid: [{}]\tmean_loss: {}\ttop1_acc: {}\tper_class_acc: {}\n'.format(
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
(epoch + 1), (mean_loss_val / val_epoch), np.mean(top1_acc_sum), 100*(correct/total)))
#############################################################################
# save model
#############################################################################
if ((epoch+1)%5==0):
checkpoint_file = os.path.join(cfg['OUT_DIR'], 'checkpoint/')
checkpoint_filename = cp.save_checkpoint(model, optimizer, epoch, gpu_nums, checkpoint_file)
log.write('[{}]--save checkpoint: {}\n'.format(
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
checkpoint_filename
))
log.close()
