def train(model, train_loader, device, optimizer, epoch, train_dict, logger):
model.train()
train_loss = 0
train_dice0 = 0
train_dice1 = 0
train_dice2 = 0
train_dice3 = 0
for batch_idx, (data, target) in enumerate(train_loader):
data = torch.squeeze(data, dim=0)
target = torch.squeeze(target, dim=0)
data, target = data.float(), target.float()
data, target = data.to(device), target.to(device)
output = model(data)
optimizer.zero_grad()
# loss = nn.CrossEntropyLoss()(output,target)
#loss=metrics.SoftDiceLoss()(output,target)
# loss=nn.MSELoss()(output,target)
loss = metrics.DiceMeanLoss()(output, target)
# loss=metrics.WeightDiceLoss()(output,target)
# loss=metrics.CrossEntropy()(output,target)
loss.backward()
optimizer.step()
train_loss = loss
train_dice0 = metrics.dice(output, target, 0)
train_dice1 = metrics.dice(output, target, 1)
train_dice2 = metrics.dice(output, target, 2)
train_dice3 = metrics.dice(output, target, 3)
print(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tdice0: {:.6f}\tdice1: {:.6f}\tdice2: {:.6f}\tdice3: {:.6f}\tT: {:.6f}\tP: {:.6f}\tTP: {:.6f}'
.format(epoch, batch_idx,
len(train_loader), 100. * batch_idx / len(train_loader),
loss.item(), train_dice0, train_dice1, train_dice2,
train_dice3, metrics.T(output, target),
metrics.P(output, target), metrics.TP(output, target)))
train_dict['loss'].append(float(train_loss))
train_dict['dice0'].append(float(train_dice0))
train_dict['dice1'].append(float(train_dice1))
train_dict['dice2'].append(float(train_dice2))
train_dict['dice3'].append(float(train_dice3))
logger.scalar_summary('train_loss', train_loss, epoch)
logger.scalar_summary('train_dice0', train_dice0, epoch)
logger.scalar_summary('train_dice1', train_dice1, epoch)
logger.scalar_summary('train_dice2', train_dice2, epoch)
logger.scalar_summary('train_dice3', train_dice3, epoch)
def val(model, val_loader, device, epoch, val_dict, logger):
model.eval()
val_loss = 0
val_dice0 = 0
val_dice1 = 0
val_dice2 = 0
val_dice3 = 0
with torch.no_grad():
for data, target in val_loader:
data = torch.squeeze(data, dim=0)
target = torch.squeeze(target, dim=0)
data, target = data.float(), target.float()
data, target = data.to(device), target.to(device)
output = model(data)
loss = metrics.DiceMeanLoss()(output, target)
#loss = metrics.SoftDiceLoss()(output, target)
dice0 = metrics.dice(output, target, 0)
dice1 = metrics.dice(output, target, 1)
dice2 = metrics.dice(output, target, 2)
dice3 = metrics.dice(output, target, 3)
val_loss += float(loss)
val_dice0 += float(dice0)
val_dice1 += float(dice1)
val_dice2 += float(dice2)
val_dice3 += float(dice3)
val_loss /= len(val_loader)
val_dice0 /= len(val_loader)
val_dice1 /= len(val_loader)
val_dice2 /= len(val_loader)
val_dice3 /= len(val_loader)
val_dict['loss'].append(float(val_loss))
val_dict['dice0'].append(float(val_dice0))
val_dict['dice1'].append(float(val_dice1))
val_dict['dice2'].append(float(val_dice2))
val_dict['dice3'].append(float(val_dice3))
logger.scalar_summary('val_loss', val_loss, epoch)
logger.scalar_summary('val_dice0', val_dice0, epoch)
logger.scalar_summary('val_dice1', val_dice1, epoch)
logger.scalar_summary('val_dice2', val_dice2, epoch)
logger.scalar_summary('val_dice3', val_dice3, epoch)
print(
'\nVal set: Average loss: {:.6f}, dice0: {:.6f}\tdice1: {:.6f}\tdice2: {:.6f}\tdice3: {:.6f}\t\n'
.format(val_loss, val_dice0, val_dice1, val_dice2, val_dice3))
logger = logger.Logger(args.log_dir)
# GENERATOR MLE TRAINING
gen_optimizer = optim.Adam(gen.parameters(), lr=1e-2)
if args.pre_g_load is not None:
print("Load pretrained MLE gen")
gen.load_state_dict(torch.load(args.pre_g_load))
else:
print('Starting Generator MLE Training...')
for epoch in range(args.mle_epochs):
print('epoch %d : ' % (epoch + 1), end='')
sys.stdout.flush()
oracle_loss, total_loss = train_generator_MLE(gen, gen_optimizer,
oracle, oracle_samples,
args)
logger.scalar_summary("oracle_loss", oracle_loss, epoch + 1)
print(' average_train_NLL = %.4f, oracle_sample_NLL = %.4f' %
(total_loss, oracle_loss))
if args.pre_g_save is not None:
torch.save(gen.state_dict(), args.pre_g_save)
# PRETRAIN DISCRIMINATOR
dis_optimizer = optim.Adagrad(dis.parameters())
if args.pre_d_load is not None:
print("Load pretrained D")
dis.load_state_dict(torch.load(args.pre_d_load))
else:
print('\nStarting Discriminator Training...')
train_discriminator(dis, dis_optimizer, oracle_samples, gen, oracle,
args.d_pre_steps, args.d_pre_epochs, args)
def validation(valid_loader, model, criterion, logger, epoch_num):
"""
Args:
train_loader:
model:
criterion:
optimizer:
epoch:
Returns:
"""
# logging accuracy and loss
valid_acc = metrics.MetricTracker()
valid_loss = metrics.MetricTracker()
valid_IoU = metrics.MetricTracker()
valid_BCE = metrics.MetricTracker()
valid_DICE = metrics.MetricTracker()
log_iter = len(valid_loader) // logger.print_freq
# switch to evaluate mode
model.eval()
# Iterate over data.
for idx, data in enumerate(tqdm(valid_loader, desc='validation')):
# get the inputs and wrap in Variable
if torch.cuda.is_available():
inputs = Variable(data['sat_img'].cuda(), volatile=True)
labels = Variable(data['map_img'].cuda(), volatile=True)
else:
inputs = Variable(data['sat_img'], volatile=True)
labels = Variable(data['map_img'], volatile=True)
# forward
# prob_map = model(inputs) # last activation was a sigmoid
# outputs = (prob_map > 0.3).float()
outputs = model(inputs)
# pay attention to the weighted loss should input logits not probs
if args.lovasz_loss:
loss, BCE_loss, DICE_loss = criterion(outputs, labels)
outputs = torch.nn.functional.sigmoid(outputs)
else:
outputs = torch.nn.functional.sigmoid(outputs)
loss, BCE_loss, DICE_loss = criterion(outputs, labels)
valid_acc.update(metrics.dice_coeff(outputs, labels), outputs.size(0))
valid_loss.update(loss.data[0], outputs.size(0))
valid_IoU.update(metrics.jaccard_index(outputs, labels),
outputs.size(0))
valid_BCE.update(BCE_loss.data[0], outputs.size(0))
valid_DICE.update(DICE_loss.data[0], outputs.size(0))
# tensorboard logging
if idx % log_iter == 0:
step = (epoch_num * logger.print_freq) + (idx / log_iter)
# log accuracy and loss
info = {
'loss': valid_loss.avg,
'accuracy': valid_acc.avg,
'IoU': valid_IoU.avg
}
for tag, value in info.items():
logger.scalar_summary(tag, value, step)
# log the sample images
log_img = [
data_utils.show_tensorboard_image(data['sat_img'],
data['map_img'],
outputs,
as_numpy=True),
]
logger.image_summary('valid_images', log_img, step)
print(
'Validation Loss: {:.4f} BCE: {:.4f} DICE: {:.4f} Acc: {:.4f} IoU: {:.4f}'
.format(valid_loss.avg, valid_BCE.avg, valid_DICE.avg, valid_acc.avg,
valid_IoU.avg))
print()
return {
'valid_loss': valid_loss.avg,
'valid_acc': valid_acc.avg,
'valid_IoU': valid_IoU.avg,
'valid_BCE': valid_BCE.avg,
'valid_DICE': valid_DICE.avg
}
def train(train_loader, model, criterion, optimizer, scheduler, logger,
epoch_num):
# logging accuracy and loss
train_acc = metrics.MetricTracker()
train_loss = metrics.MetricTracker()
train_IoU = metrics.MetricTracker()
train_BCE = metrics.MetricTracker()
train_DICE = metrics.MetricTracker()
meters = {
"train_acc": train_acc,
"train_loss": train_loss,
"train_IoU": train_IoU,
"train_BCE": train_BCE,
"train_DICE": train_DICE,
"outputs": None
}
log_iter = len(train_loader) // logger.print_freq
model.train()
scheduler.step()
cache = None
cached_loss = 0
# iterate over data
for idx, data in enumerate(tqdm(train_loader, desc="training")):
meters = make_train_step(idx, data, model, optimizer, criterion,
meters)
# hard negative mining
# https://github.com/asanakoy/kaggle_carvana_segmentation/blob/master/albu/src/train.py
if args.hard_mining:
if cache is None or cached_loss < meters["train_loss"].val:
cached_loss = meters["train_loss"].val
cache = deepcopy(data)
if idx % 50 == 0 and cache is not None:
meters = make_train_step(idx, data, model, optimizer,
criterion, meters)
cache = None
cached_loss = 0
# tensorboard logging
if idx % log_iter == 0:
step = (epoch_num * logger.print_freq) + (idx / log_iter)
# log accuracy and loss
info = {
'loss': meters["train_loss"].avg,
'accuracy': meters["train_acc"].avg,
'IoU': meters["train_IoU"].avg
}
for tag, value in info.items():
logger.scalar_summary(tag, value, step)
# # log weights, biases, and gradients
# for tag, value in model.named_parameters():
# tag = tag.replace('.', '/')
# logger.histo_summary(tag, value.data.cpu().numpy(), step)
# logger.histo_summary(tag + '/grad', value.grad.data.cpu().numpy(), step)
# log the sample images
log_img = [
data_utils.show_tensorboard_image(data['sat_img'],
data['map_img'],
meters["outputs"],
as_numpy=True),
]
logger.image_summary('train_images', log_img, step)
print(
'Training Loss: {:.4f} BCE: {:.4f} DICE: {:.4f} Acc: {:.4f} IoU: {:.4f} '
.format(meters["train_loss"].avg, meters["train_BCE"].avg,
meters["train_DICE"].avg, meters["train_acc"].avg,
meters["train_IoU"].avg))
print()
return {
'train_loss': meters["train_loss"].avg,
'train_acc': meters["train_acc"].avg,
'train_IoU': meters["train_IoU"].avg,
'train_BCE': meters["train_BCE"].avg,
'train_DICE': meters["train_DICE"].avg
}
step, blobs['im_name'], train_loss / step_cnt, fps, 1./fps, lr, momentum, weight_decay)
log_print(log_text, color='green', attrs=['bold'])
re_cnt = True
#TODO: evaluate the model every 5000 iterations (N defined in handout)
if step%100 == 0 and step>0:
save_name = 'test_'+str(step)
net.eval()
imdb_test.competition_mode(on=True)
ap = test_net(save_name, net, imdb_test,
cfg.TRAIN.BATCH_SIZE, thresh=thresh, visualize=use_visdom)
# # import pdb; pdb.set_trace()
for i in range(len(imdb.classes)):
cls_name = imdb.classes[i]
cur_ap = ap[i]
logger.scalar_summary('{}_AP'.format(cls_name), cur_ap, step)
viz.updateTrace(X=np.array([step]), Y=np.array([np.average(ap)]), win=ap_win,name='train_ap')
#TODO: Perform all visualizations here
#You can define other interval variable if you want (this is just an
#example)
#The intervals for different things are defined in the handout
# if visualize and step%vis_interval==0:
# #TODO: Create required visualizations
# print("visualizing")
# if use_tensorboard:
# print('Logging to Tensorboard')
# data_log.scalar_summary(tag='train/loss', value=loss, step=step)
outputs = cnn.forward(vectors_in)
# evaluate
loss = lossF(outputs, ans)
loss.backward()
optimizer.step()
optimizer.zero_grad()
corr = np.sum(
np.argmax(outputs.detach().numpy(), axis=1) ==
ans.detach().numpy())
train_accu.append(corr / BS)
logger.scalar_summary("loss", loss, (N_train // BS) *
(DATE_LEN - INPUT_LEN - 1) * epo +
N_train // BS * segment_idx + idx)
logger.scalar_summary("accu", corr / BS, (N_train // BS) *
(DATE_LEN - INPUT_LEN - 1) * epo +
N_train // BS * segment_idx + idx)
print("segment idx = {}".format(segment_idx))
print("train: epo = {}, train-accu = {}".format(
epo, np.mean(train_accu)))
# validation
cnn.eval()
valid_accu = []
for segment_idx in range(DATE_LEN - INPUT_LEN - 1):
vectors_in = torch.Tensor(validset[:, 1:, segment_idx:segment_idx +