def train(resume=False):
it = 0
writer = SummaryWriter('../runs/' + hparams.exp_name)
for k in hparams.__dict__.keys():
writer.add_text(str(k), str(hparams.__dict__[k]))
train_dataset = ChestData(
data_csv=hparams.train_csv,
data_dir=hparams.train_dir,
transform=transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize((0.485), (0.229))
]))
validation_dataset = ChestData(
data_csv=hparams.valid_csv,
data_dir=hparams.valid_dir,
transform=transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize((0.485), (0.229))
]))
train_loader = DataLoader(train_dataset,
batch_size=hparams.batch_size,
shuffle=True,
num_workers=0)
validation_loader = DataLoader(validation_dataset,
batch_size=hparams.batch_size,
shuffle=True,
num_workers=0)
print('loaded train data of length : {}'.format(len(train_dataset)))
Tensor = torch.cuda.FloatTensor if hparams.cuda else torch.FloatTensor
def validation(encoder_, decoder_=None, send_stats=False, epoch=0):
encoder_ = encoder_.eval()
if decoder_:
decoder_ = decoder_.eval()
# print('Validating model on {0} examples. '.format(len(validation_loader)))
with torch.no_grad():
scores_list = []
labels_list = []
val_loss = 0
for (img, labels, imgs_names) in validation_loader:
img = Variable(img.float(), requires_grad=False)
labels = Variable(labels.float(), requires_grad=False)
scores = None
if hparams.cuda:
img = img.cuda(hparams.gpu_device)
labels = labels.cuda(hparams.gpu_device)
z = encoder_(img)
if decoder_:
outputs = decoder_(z)
scores = torch.sum(
(outputs - img)**2, dim=tuple(range(
1, outputs.dim()))) # (outputs - img) ** 2
# rec_loss = rec_loss.view(outputs.shape[0], -1)
# rec_loss = torch.sum(torch.sum(rec_loss, dim=1))
val_loss += torch.sum(scores)
save_image(img,
'tmp/img_{}.png'.format(epoch),
normalize=True)
save_image(outputs,
'tmp/reconstructed_{}.png'.format(epoch),
normalize=True)
else:
dist = torch.sum((z - encoder.center)**2, dim=1)
if hparams.objective == 'soft-boundary':
scores = dist - encoder.radius**2
val_loss += (1 / hparams.nu) * torch.sum(
torch.max(torch.zeros_like(scores), scores))
else:
scores = dist
val_loss += torch.sum(dist)
scores_list.append(scores)
labels_list.append(labels)
scores = torch.cat(scores_list, dim=0)
labels = torch.cat(labels_list, dim=0)
val_loss /= len(validation_dataset)
val_loss += encoder_.radius**2 if decoder_ and hparams.objective == 'soft-boundary' else 0
if hparams.cuda:
labels = labels.cpu()
scores = scores.cpu()
labels = labels.view(-1).numpy()
scores = scores.view(-1).detach().numpy()
auc = roc_auc_score(labels, scores)
return auc, val_loss
### validation function ends.
if hparams.cuda:
encoder = Encoder().cuda(hparams.gpu_device)
decoder = Decoder().cuda(hparams.gpu_device)
else:
encoder = Encoder()
decoder = Decoder()
params_count = 0
for param in encoder.parameters():
params_count += np.prod(param.size())
for param in decoder.parameters():
params_count += np.prod(param.size())
print('Model has {0} trainable parameters'.format(params_count))
if not hparams.load_model:
encoder.apply(weights_init_normal)
decoder.apply(weights_init_normal)
optim_params = list(encoder.parameters())
optimizer_train = optim.Adam(optim_params,
lr=hparams.train_lr,
weight_decay=hparams.weight_decay,
amsgrad=hparams.optimizer == 'amsgrad')
if hparams.pretrain:
optim_params += list(decoder.parameters())
optimizer_pre = optim.Adam(optim_params,
lr=hparams.pretrain_lr,
weight_decay=hparams.ae_weight_decay,
amsgrad=hparams.optimizer == 'amsgrad')
# scheduler_pre = ReduceLROnPlateau(optimizer_pre, mode='min', factor=0.5, patience=10, verbose=True, cooldown=20)
scheduler_pre = MultiStepLR(optimizer_pre,
milestones=hparams.lr_milestones,
gamma=0.1)
# scheduler_train = ReduceLROnPlateau(optimizer_train, mode='min', factor=0.5, patience=10, verbose=True, cooldown=20)
scheduler_train = MultiStepLR(optimizer_train,
milestones=hparams.lr_milestones,
gamma=0.1)
print('Starting training.. (log saved in:{})'.format(hparams.exp_name))
start_time = time.time()
mode = 'pretrain' if hparams.pretrain else 'train'
best_valid_loss = 100000000000000000
best_valid_auc = 0
encoder = init_center(encoder, train_loader)
# print(model)
for epoch in range(hparams.num_epochs):
if mode == 'pretrain' and epoch == hparams.pretrain_epoch:
print('Pretraining done.')
mode = 'train'
best_valid_loss = 100000000000000000
best_valid_auc = 0
encoder = init_center(encoder, train_loader)
for batch, (imgs, labels, _) in enumerate(train_loader):
# imgs = Variable(imgs.float(), requires_grad=False)
if hparams.cuda:
imgs = imgs.cuda(hparams.gpu_device)
if mode == 'pretrain':
optimizer_pre.zero_grad()
z = encoder(imgs)
outputs = decoder(z)
# print(torch.max(outputs), torch.mean(imgs), torch.min(outputs), torch.mean(imgs))
scores = torch.sum((outputs - imgs)**2,
dim=tuple(range(1, outputs.dim())))
# print(scores)
loss = torch.mean(scores)
loss.backward()
optimizer_pre.step()
writer.add_scalar('pretrain_loss',
loss.item(),
global_step=batch +
len(train_loader) * epoch)
else:
optimizer_train.zero_grad()
z = encoder(imgs)
dist = torch.sum((z - encoder.center)**2, dim=1)
if hparams.objective == 'soft-boundary':
scores = dist - encoder.radius**2
loss = encoder.radius**2 + (1 / hparams.nu) * torch.mean(
torch.max(torch.zeros_like(scores), scores))
else:
loss = torch.mean(dist)
loss.backward()
optimizer_train.step()
if hparams.objective == 'soft-boundary' and epoch >= hparams.warmup_epochs:
R = np.quantile(np.sqrt(dist.clone().data.cpu().numpy()),
1 - hparams.nu)
encoder.radius = torch.tensor(R)
if hparams.cuda:
encoder.radius = encoder.radius.cuda(
hparams.gpu_device)
writer.add_scalar('radius',
encoder.radius.item(),
global_step=batch +
len(train_loader) * epoch)
writer.add_scalar('train_loss',
loss.item(),
global_step=batch +
len(train_loader) * epoch)
# pred_labels = (scores >= hparams.thresh)
# save_image(imgs, 'train_imgs.png')
# save_image(noisy_imgs, 'train_noisy.png')
# save_image(gen_imgs, 'train_z.png')
if batch % hparams.print_interval == 0:
print('[Epoch - {0:.1f}, batch - {1:.3f}, loss - {2:.6f}]'.\
format(1.0*epoch, 100.0*batch/len(train_loader), loss.item()))
if mode == 'pretrain':
val_auc, rec_loss = validation(copy.deepcopy(encoder),
copy.deepcopy(decoder),
epoch=epoch)
else:
val_auc, val_loss = validation(copy.deepcopy(encoder), epoch=epoch)
writer.add_scalar('val_auc', val_auc, global_step=epoch)
if mode == 'pretrain':
best_valid_auc = max(best_valid_auc, val_auc)
scheduler_pre.step()
writer.add_scalar('rec_loss', rec_loss, global_step=epoch)
writer.add_scalar('pretrain_lr',
optimizer_pre.param_groups[0]['lr'],
global_step=epoch)
torch.save(
{
'epoch': epoch,
'encoder_state_dict': encoder.state_dict(),
'decoder_state_dict': decoder.state_dict(),
'optimizer_pre_state_dict': optimizer_pre.state_dict(),
}, hparams.model + '.pre')
if best_valid_loss >= rec_loss:
best_valid_loss = rec_loss
torch.save(
{
'epoch': epoch,
'encoder_state_dict': encoder.state_dict(),
'decoder_state_dict': decoder.state_dict(),
'optimizer_pre_state_dict': optimizer_pre.state_dict(),
}, hparams.model + '.pre.best')
print('best model on validation set saved.')
print('[Epoch - {0:.1f} ---> rec_loss - {1:.4f}, current_lr - {2:.6f}, val_auc - {3:.4f}, best_valid_auc - {4:.4f}] - time - {5:.1f}'\
.format(1.0*epoch, rec_loss, optimizer_pre.param_groups[0]['lr'], val_auc, best_valid_auc, time.time()-start_time))
else:
scheduler_train.step()
writer.add_scalar('val_loss', val_loss, global_step=epoch)
writer.add_scalar('train_lr',
optimizer_train.param_groups[0]['lr'],
global_step=epoch)
torch.save(
{
'epoch': epoch,
'encoder_state_dict': encoder.state_dict(),
'center': encoder.center,
'radius': encoder.radius,
'optimizer_train_state_dict': optimizer_train.state_dict(),
}, hparams.model + '.train')
if best_valid_loss >= val_loss:
best_valid_loss = val_loss
torch.save(
{
'epoch': epoch,
'encoder_state_dict': encoder.state_dict(),
'center': encoder.center,
'radius': encoder.radius,
'optimizer_train_state_dict':
optimizer_train.state_dict(),
}, hparams.model + '.train.best')
print('best model on validation set saved.')
if best_valid_auc <= val_auc:
best_valid_auc = val_auc
torch.save(
{
'epoch': epoch,
'encoder_state_dict': encoder.state_dict(),
'center': encoder.center,
'radius': encoder.radius,
'optimizer_train_state_dict':
optimizer_train.state_dict(),
}, hparams.model + '.train.auc')
print('best model on validation set saved.')
print('[Epoch - {0:.1f} ---> val_loss - {1:.4f}, current_lr - {2:.6f}, val_auc - {3:.4f}, best_valid_auc - {4:.4f}] - time - {5:.1f}'\
.format(1.0*epoch, val_loss, optimizer_train.param_groups[0]['lr'], val_auc, best_valid_auc, time.time()-start_time))
start_time = time.time()
def test(model_path,
send_stats=False,
accuracy=False,
data='valid',
plot_name='valid'):
if data == 'valid':
test_dataset = ChestData(
data_csv=hparams.valid_csv,
data_dir=hparams.valid_dir,
transform=transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize((0.485), (0.229))
]))
else:
test_dataset = ChestData(
data_csv=hparams.test_csv,
data_dir=hparams.test_dir,
transform=transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize((0.485), (0.229))
]))
test_loader = DataLoader(test_dataset,
batch_size=hparams.batch_size,
shuffle=True,
num_workers=2)
if hparams.cuda:
encoder = Encoder().cuda(hparams.gpu_device)
checkpoint = torch.load(model_path, map_location=hparams.gpu_device)
encoder.load_state_dict(checkpoint['encoder_state_dict'])
encoder.center = checkpoint['center']
encoder.radius = checkpoint['radius']
else:
encoder = Encoder()
checkpoint = torch.load(model_path, map_location='cpu')
encoder.load_state_dict(checkpoint['encoder_state_dict'])
encoder.center = checkpoint['center']
encoder.radius = checkpoint['radius']
# print('Model loaded')
Tensor = torch.cuda.FloatTensor if hparams.cuda else torch.FloatTensor
print('Testing model on {0} examples. '.format(len(test_dataset)))
encoder = encoder.eval()
# print('Validating model on {0} examples. '.format(len(validation_loader)))
# print(encoder.center)
# print(encoder.radius)
with torch.no_grad():
scores_list = []
labels_list = []
test_loss = 0
for (img, labels, _) in tqdm(test_loader):
# img = Variable(img.float(), requires_grad=False)
# labels = Variable(labels.float(), requires_grad=False)
scores = None
if hparams.cuda:
img = img.cuda(hparams.gpu_device)
labels = labels.cuda(hparams.gpu_device)
z = encoder(img)
dist = torch.sum((z - encoder.center)**2, dim=1)
if hparams.objective == 'soft-boundary':
scores = dist - encoder.radius**2
test_loss += (1 / hparams.nu) * torch.sum(
torch.max(torch.zeros_like(scores), scores))
else:
scores = dist
test_loss += torch.sum(dist)
scores_list.append(scores)
labels_list.append(labels)
scores = torch.cat(scores_list, dim=0)
labels = torch.cat(labels_list, dim=0)
if hparams.cuda:
plot_labels = labels.cpu()
plot_scores = scores.cpu()
plot_labels = plot_labels.numpy()
plot_scores = plot_scores.numpy()
accuracy = -1
best_thr = -1
if plot_name:
groups = ("inlier", "outlier")
color = ['red' if l == 0 else 'green' for l in plot_labels]
lbl_name = ['inlier' if l == 0 else 'outlier' for l in plot_labels]
plt.scatter(plot_scores[plot_labels == 1],
plot_labels[plot_labels == 1],
color='red',
edgecolors='none',
label='outlier',
marker='+')
plt.scatter(plot_scores[plot_labels == 0],
plot_labels[plot_labels == 0],
color='green',
edgecolors='none',
label='inlier',
marker='+')
plt.title('Scores vs labels')
plt.legend(loc=0)
plt.savefig(hparams.result_dir + plot_name + '_scores.png')
print('scores plot saved to {}'.format(hparams.result_dir +
plot_name + '_scores.png'))
# Compute ROC curve and ROC area for each class
# plot_scores -= np.min(plot_scores)
# plot_scores /= np.max(plot_scores)
fpr = dict()
tpr = dict()
roc_auc = dict()
for i in range(1):
fpr[i], tpr[i], thresholds = sklearn.metrics.roc_curve(
plot_labels, plot_scores)
roc_auc[i] = sklearn.metrics.auc(fpr[i], tpr[i])
# Compute micro-average ROC curve and ROC area
fpr["micro"], tpr["micro"], _ = sklearn.metrics.roc_curve(
plot_labels.ravel(), plot_scores.ravel())
roc_auc["micro"] = sklearn.metrics.auc(fpr["micro"], tpr["micro"])
plt.figure()
lw = 2
plt.plot(fpr[0],
tpr[0],
color='darkorange',
lw=lw,
label='ROC curve (area = %0.2f)' % roc_auc[0])
plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title(
'Receiver operating characteristic example for {} set.'.format(
plot_name))
plt.legend(loc="lower right")
plt.savefig(hparams.result_dir + plot_name + '_roc_curve.png')
print(
'roc-auc curve saved to {}'.format(hparams.result_dir +
plot_name + '_scores.png'))
for thr in thresholds:
pred_labels = 1.0 * np.array(plot_scores >= thr)
acc = np.mean(1.0 * (pred_labels == plot_labels))
if accuracy <= acc:
best_thr = thr
accuracy = acc
test_loss /= len(test_dataset)
test_loss += encoder.radius**2 if hparams.objective == 'soft-boundary' else 0
if hparams.cuda:
labels = labels.cpu()
scores = scores.cpu()
labels = labels.view(-1).numpy()
scores = scores.view(-1).detach().numpy()
auc = roc_auc_score(labels, scores)
print(
'== Test on -- ' + model_path +
' == auc - {0:.4f}, test_loss - {1:.4f}, best_threshold - {2:.4f}, accuracy - {3:.4f} =='
.format(auc, test_loss, best_thr, accuracy))
return auc, test_loss
