def train(model, train_loader, val_loader, optimizer, num_epochs,
path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1).to(device) # Use for NLLLoss()
softmax = nn.Softmax(dim=1).to(device)
criterion_nlloss = nn.NLLLoss(size_average=False).to(device)
criterion_mseloss = nn.MSELoss(size_average=False).to(device)
metrics_evaluator = PerformanceMetricsEvaluator()
to_tensor = transforms.ToTensor()
writer = SummaryWriter('runs/shape_net_espnet_with_pretraining/')
since = time.time()
best_model_weights = model.state_dict()
best_IoU = 0.0
best_val_loss = 1000000000
curr_val_loss = 0.0
curr_training_loss = 0.0
curr_training_IoU = 0.0
curr_val_IoU = 0.0
curr_unet_training_IoU = 0.0
curr_unet_val_IoU = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step(best_val_loss)
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
running_IoU = 0.0
unet_IoU = 0.0
# Iterate over data.
for imgs, masks in tqdm(data_loader):
mask_to_encode = (
np.arange(2) == masks.numpy()[..., None]).astype(float)
mask_to_encode = torch.from_numpy(
np.moveaxis(mask_to_encode, 3, 1)).float().to(device)
imgs, masks = imgs.to(device), masks.to(device)
# masks_for_shape = masks.clone().unsqueeze(1).float()
# zero the parameter gradients
optimizer.zero_grad()
# forward
if phase == 'train':
unet_prediction, unet_prediction_1, shape_net_encoded_prediction, shape_net_final_prediction = model(
imgs)
else:
unet_prediction, shape_net_encoded_prediction, shape_net_final_prediction = model(
imgs)
encoded_mask = model(mask_to_encode, only_encode=True)
log_softmax_unet_prediction = log_softmax(unet_prediction)
if phase == 'train':
log_softmax_unet_prediction_1 = log_softmax(
unet_prediction_1)
third_term = criterion_nlloss(log_softmax_unet_prediction, masks) + \
criterion_nlloss(log_softmax_unet_prediction_1, masks)
else:
third_term = criterion_nlloss(log_softmax_unet_prediction,
masks)
softmax_unet_prediction = softmax(unet_prediction)
softmax_shape_net_final_prediction = softmax(
shape_net_final_prediction)
log_softmax_unet_prediction = log_softmax(unet_prediction)
# log_softmax_shape_net_final_prediction = log_softmax(shape_net_final_prediction)
# first_term = criterion_nlloss(log_softmax_unet_prediction, log_softmax_shape_net_final_prediction)
first_term = criterion_mseloss(
softmax_unet_prediction,
softmax_shape_net_final_prediction)
second_term = criterion_mseloss(encoded_mask,
shape_net_encoded_prediction)
# third_term = criterion_mseloss(softmax_unet_prediction, mask_to_encode)
print('First term: ',
first_term.data.cpu().numpy(), 'Second term: ',
second_term.data.cpu().numpy(), 'Third term: ',
third_term.data.cpu().numpy())
# print()
lambda_1 = 0.5
lambda_2 = 0.5
loss = first_term + lambda_1 * third_term + lambda_2 * third_term
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
collapsed_softmax_logits = np.argmax(
softmax_shape_net_final_prediction.detach(), axis=1)
collapsed_softmax_unet = np.argmax(
softmax_unet_prediction.detach(), axis=1)
prediction = np.argmax(softmax_shape_net_final_prediction[
0, :, :, :].detach().cpu().numpy(),
axis=0)
prediction_segm_net = np.argmax(
softmax_unet_prediction[0, :, :, :].detach().cpu().numpy(),
axis=0)
empty_channel = np.zeros((1024, 1024), dtype=np.uint64)
# final_vis_img = np.stack([prediction, masks[0], empty_channel], axis=0)
_, mask_contours, _ = cv2.findContours(
np.expand_dims(masks[0].cpu().numpy().astype(np.uint8),
axis=-1), cv2.RETR_TREE,
cv2.CHAIN_APPROX_NONE)
_, pred_contours, _ = cv2.findContours(
prediction.astype(np.uint8), cv2.RETR_TREE,
cv2.CHAIN_APPROX_NONE)
_, pred_segm_net_contours, _ = cv2.findContours(
prediction_segm_net.astype(np.uint8), cv2.RETR_TREE,
cv2.CHAIN_APPROX_NONE)
# pred_contours = [cnt for cnt in pred_contours if cv2.contourArea(cnt) > 1000]
res_img = np.copy(imgs[0])
fof = (np.moveaxis(res_img, 0, -1) * 255).astype(
np.uint8).copy()
cv2.drawContours(fof, mask_contours, -1, (0, 255, 0), 2)
cv2.drawContours(fof, pred_contours, -1, (255, 0, 0), 2)
cv2.drawContours(fof, pred_segm_net_contours, -1, (0, 0, 255),
1)
if phase == 'val':
name = 'ValidationEpoch'
else:
name = 'TrainingEpoch'
writer.add_image('{}: {}'.format(name, str(epoch)),
np.moveaxis(fof, -1, 0), epoch)
# statistics
running_loss += loss.detach().item()
batch_IoU = 0.0
for k in range(len(imgs)):
batch_IoU += metrics_evaluator.mean_IU(
collapsed_softmax_logits.numpy()[k],
masks.cpu().numpy()[k])
batch_IoU /= len(imgs)
unet_batch_IoU = 0.0
for j in range(len(imgs)):
unet_batch_IoU += metrics_evaluator.mean_IU(
collapsed_softmax_logits.numpy()[k],
masks.cpu().numpy()[k])
unet_batch_IoU /= len(imgs)
running_IoU += batch_IoU
unet_IoU += unet_batch_IoU
epoch_loss = running_loss / len(data_loader)
epoch_IoU = running_IoU / len(data_loader)
epoch_unet_IoU = unet_IoU / len(data_loader)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss,
epoch_IoU))
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
best_val_loss = epoch_loss
best_IoU = epoch_IoU
best_model_weights = model.state_dict()
if phase == 'val':
# print(optimizer.param_groups[0]['lr'])
curr_val_loss = epoch_loss
curr_val_IoU = epoch_IoU
curr_unet_val_IoU = epoch_unet_IoU
else:
curr_training_loss = epoch_loss
curr_training_IoU = epoch_IoU
curr_unet_training_IoU = epoch_unet_IoU
writer.add_scalars(
'TrainValIoU', {
'trainIoU': curr_training_IoU,
'validationIoU': curr_val_IoU,
'trainUnetIoU': curr_unet_training_IoU,
'validationUnetIoU': curr_unet_val_IoU
}, epoch)
writer.add_scalars('TrainValLoss', {
'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
}, epoch)
# Saving best model
torch.save(
best_model_weights,
os.path.join(
path_to_save_best_weights,
'shape_net_espnet_with_pretraining{:2f}.pth'.format(
best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
def train(model, train_loader, val_loader, optimizer, num_epochs, scheduler,
path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1) # Use for NLLLoss()
softmax = nn.Softmax(dim=1)
criterion_nlloss = nn.NLLLoss().to(device)
criterion_mseloss = nn.MSELoss().to(device)
metrics_evaluator = PerformanceMetricsEvaluator()
writer = SummaryWriter('FrameworkTensorboard/rabish/')
since = time.time()
best_model_weights = model.state_dict()
best_IoU = 0.0
best_val_loss = 1000000000
curr_val_loss = 0.0
curr_training_loss = 0.0
curr_training_IoU = 0.0
curr_val_IoU = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step(best_val_loss)
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
running_IoU = 0
# Iterate over data.
for imgs, masks in tqdm(data_loader):
mask_to_encode = masks.numpy()
mask_to_encode = (
np.arange(4) == mask_to_encode[..., None]).astype(float)
mask_to_encode = torch.from_numpy(
np.moveaxis(mask_to_encode, 3, 1)).float().to(device)
imgs, masks = imgs.to(device), masks.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
unet_prediction, shape_net_encoded_prediction, shape_net_final_prediction = model(
imgs)
encoded_mask = model(mask_to_encode, only_encode=True)
log_softmax_unet_prediction = log_softmax(unet_prediction)
softmax_unet_prediction = softmax(unet_prediction)
softmax_shape_net_final_prediction = softmax(
shape_net_final_prediction)
np.save('unet_prediction',
unet_prediction.cpu().detach().numpy())
np.save('shapeNet_prediction',
shape_net_final_prediction.cpu().detach().numpy())
first_term = criterion_mseloss(unet_prediction.detach(),
shape_net_final_prediction)
second_term = criterion_mseloss(encoded_mask,
shape_net_encoded_prediction)
third_term = criterion_nlloss(log_softmax_unet_prediction,
masks)
# third_term = criterion_nlloss(log_softmax_unet_prediction, masks)
lambda_1 = 0.5
lambda_2 = 0.5
lambda_3 = 100
loss = first_term + lambda_1 * third_term + lambda_2 * third_term
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
unet_softmax_collupsed = np.argmax(
softmax_unet_prediction.detach(), axis=1)
unet_softmax_uncollupsed = (np.arange(
4) == unet_softmax_collupsed[..., None]).numpy().astype(
np.uint8)
unet_softmax_uncollupsed = torch.from_numpy(
np.moveaxis(unet_softmax_uncollupsed, 3, 1)).float()
softmax_shape_net_final_collupsed = np.argmax(
softmax_shape_net_final_prediction.detach(), axis=1)
softmax_shape_net_final_uncollupsed = (
np.arange(4) == softmax_shape_net_final_collupsed[...,
None]
).numpy().astype(np.uint8)
softmax_shape_net_final_uncollupsed = torch.from_numpy(
np.moveaxis(softmax_shape_net_final_uncollupsed, 3,
1)).float()
for i in range(len(unet_softmax_uncollupsed)):
empty_channel = np.zeros((544, 544), dtype=np.uint64)
if phase == 'val':
img_name = 'ValidationEpoch: {}'.format(str(epoch))
else:
img_name = 'TrainingEpoch: {}'.format(str(epoch))
writer.add_image(
img_name,
vutils.make_grid([
imgs[i].detach().cpu(),
unet_softmax_uncollupsed[i][1:, :, :],
softmax_shape_net_final_uncollupsed[i][
1:, :, :].cpu()
]), epoch)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.detach().item()
running_IoU += metrics_evaluator.mean_IU(
softmax_shape_net_final_collupsed.numpy()[0],
masks.cpu().numpy()[0])
epoch_loss = running_loss / len(data_loader)
epoch_IoU = running_IoU / len(data_loader)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss,
epoch_IoU))
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
best_val_loss = epoch_loss
best_IoU = epoch_IoU
best_model_weights = model.state_dict()
if phase == 'val':
# print(optimizer.param_groups[0]['lr'])
curr_val_loss = epoch_loss
curr_val_IoU = epoch_IoU
else:
curr_training_loss = epoch_loss
curr_training_IoU = epoch_IoU
writer.add_scalars('TrainValIoU', {
'trainIoU': curr_training_IoU,
'validationIoU': curr_val_IoU
}, epoch)
writer.add_scalars('TrainValLoss', {
'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
}, epoch)
# Saving best model
torch.save(
best_model_weights,
os.path.join(
path_to_save_best_weights,
'Unet_manually_corrupted_lumenTissue{:2f}.pth'.format(
best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
test_loader = torch.utils.data.DataLoader(shape_test, batch_size=1, shuffle=False)
# Model Creation
device = torch.device("cuda:{}".format(str(get_freer_gpu())))
PATH_TO_THE_WEIGHTS = 'weights/unet0.180934.pth'
model = UNet((3,512,512))
# model = SH_UNet()
# model = R2AttU_Net(img_ch=3, output_ch=15)
# model = SH_2UNet()
model.load_state_dict(torch.load(PATH_TO_THE_WEIGHTS))
model.to(device)
# Evaluation Techniques
metrics_evaluator = PerformanceMetricsEvaluator()
# Evaluation
iou_of_the_model = 0
mean_acc_of_the_model = 0
softmax = nn.Softmax(dim=1)
temporary_counter = 0
# with torch.no_grad():
# unique_ind = 0
# to_tensor = transforms.ToTensor()
# for (test_imgs, test_masks) in tqdm(test_loader):
def train(model, train_loader, val_loader, optimizer, num_epochs, path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1)# Use for NLLLoss()
softmax = nn.Softmax(dim=1)
criterion_nlloss = nn.NLLLoss()
metrics_evaluator = PerformanceMetricsEvaluator()
to_tensor = transforms.ToTensor()
writer = SummaryWriter('runs/BiSeNet_with_couple_losses/')
since = time.time()
best_model_weights = model.state_dict()
best_IoU = 0.0
best_val_loss = 1000000000
curr_val_loss = 0.0
curr_training_loss = 0.0
curr_training_IoU = 0.0
curr_val_IoU = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step(best_val_loss)
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
running_IoU = 0
# Iterate over data.
for imgs, masks in tqdm(data_loader):
imgs, masks = imgs.to(device), masks.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
if phase == 'train':
logits, logits_2, logits_3 = model(imgs)
log_logits = log_softmax(logits).to(device)
log_logits_2 = log_softmax(logits_2).to(device)
log_logits_3 = log_softmax(logits_3).to(device)
loss = criterion_nlloss(log_logits, masks) + criterion_nlloss(log_logits_2, masks) + criterion_nlloss(log_logits_3, masks)
else:
logits = model(imgs)
log_logits = log_logits = log_softmax(logits).to(device)
loss = criterion_nlloss(log_logits, masks)
# log_logits = log_softmax(logits).to(device)
# Save the results of the epoch
# output = np.argmax(logits.detach().cpu().numpy(), axis=1)
output = logits.detach().cpu().numpy()
# FOR SAVING PREDICTIONS
# corr_masks = masks.detach().cpu().numpy()
# for i in range(len(output)):
# predicted_log_logits.append(output[i])
# corresponding_masks.append(corr_masks[i])
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
softmax_logits = softmax(logits)
collapsed_softmax_logits = np.argmax(softmax_logits.detach(), axis=1)
prediction = np.argmax(softmax_logits[0,:, :, :].detach().cpu().numpy(), axis=0)
empty_channel = np.zeros((1024, 1024), dtype=np.uint64)
# final_vis_img = np.stack([prediction, masks[0], empty_channel], axis=0)
_, mask_contours, _ = cv2.findContours(np.expand_dims(masks[0].cpu().numpy().astype(np.uint8), axis=-1),
cv2.RETR_TREE,cv2.CHAIN_APPROX_NONE)
_, pred_contours, _ = cv2.findContours(prediction.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# pred_contours = [cnt for cnt in pred_contours if cv2.contourArea(cnt) > 1000]
res_img = np.copy(imgs[0])
fof = (np.moveaxis(res_img, 0, -1)*255).astype(np.uint8).copy()
cv2.drawContours(fof, mask_contours, -1, (0,255,0), 2)
cv2.drawContours(fof, pred_contours, -1, (255,0,0), 2)
if phase == 'val':
name = 'ValidationEpoch'
else:
name = 'TrainingEpoch'
writer.add_image('{}: {}'.format(name, str(epoch)), np.moveaxis(fof, -1, 0),epoch)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.detach().item()
batch_IoU = 0.0
for k in range(len(imgs)):
batch_IoU += metrics_evaluator.mean_IU(collapsed_softmax_logits.numpy()[k], masks.cpu().numpy()[k])
batch_IoU /= len(imgs)
running_IoU += batch_IoU
epoch_loss = running_loss / len(data_loader)
epoch_IoU = running_IoU / len(data_loader)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_IoU))
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
best_val_loss = epoch_loss
best_IoU = epoch_IoU
best_model_weights = model.state_dict()
if phase == 'val':
# print(optimizer.param_groups[0]['lr'])
curr_val_loss = epoch_loss
curr_val_IoU = epoch_IoU
else:
curr_training_loss = epoch_loss
curr_training_IoU = epoch_IoU
writer.add_scalars('TrainValIoU',
{'trainIoU': curr_training_IoU,
'validationIoU': curr_val_IoU
},
epoch
)
writer.add_scalars('TrainValLoss',
{'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
},
epoch
)
# Saving best model
torch.save(best_model_weights,
os.path.join(path_to_save_best_weights, 'BiSeNet_with_couple_losses{:2f}.pth'.format(best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
def train(model, train_loader, val_loader, optimizer, num_epochs,
path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1).to(device) # Use for NLLLoss()
softmax = nn.Softmax(dim=1).to(device)
criterion_nlloss = nn.NLLLoss(size_average=False).to(device)
criterion_mseloss = nn.MSELoss(size_average=False).to(device)
metrics_evaluator = PerformanceMetricsEvaluator()
to_tensor = transforms.ToTensor()
writer = SummaryWriter('runs/shape_net_with_pretraining/')
since = time.time()
best_model_weights = model.state_dict()
best_IoU = 0.0
best_val_loss = 1000000000
curr_val_loss = 0.0
curr_training_loss = 0.0
curr_training_IoU = 0.0
curr_val_IoU = 0.0
curr_unet_training_IoU = 0.0
curr_unet_val_IoU = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step(best_val_loss)
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
running_IoU = 0.0
unet_IoU = 0.0
# Iterate over data.
for imgs, masks in tqdm(data_loader):
mask_to_encode = (
np.arange(15) == masks.numpy()[..., None]).astype(float)
mask_to_encode = torch.from_numpy(
np.moveaxis(mask_to_encode, 3, 1)).float().to(device)
imgs, masks = imgs.to(device), masks.to(device)
# masks_for_shape = masks.clone().unsqueeze(1).float()
# zero the parameter gradients
optimizer.zero_grad()
# forward
unet_prediction, shape_net_encoded_prediction, shape_net_final_prediction = model(
imgs)
encoded_mask = model(mask_to_encode, only_encode=True)
log_softmax_unet_prediction = log_softmax(unet_prediction)
softmax_unet_prediction = softmax(unet_prediction)
softmax_shape_net_final_prediction = softmax(
shape_net_final_prediction)
log_softmax_unet_prediction = log_softmax(unet_prediction)
# log_softmax_shape_net_final_prediction = log_softmax(shape_net_final_prediction)
# first_term = criterion_nlloss(log_softmax_unet_prediction, log_softmax_shape_net_final_prediction)
first_term = criterion_mseloss(
softmax_unet_prediction,
softmax_shape_net_final_prediction)
second_term = criterion_mseloss(encoded_mask,
shape_net_encoded_prediction)
third_term = criterion_nlloss(log_softmax_unet_prediction,
masks)
# third_term = criterion_mseloss(softmax_unet_prediction, mask_to_encode)
print('First term: ',
first_term.data.cpu().numpy(), 'Second term: ',
second_term.data.cpu().numpy(), 'Third term: ',
third_term.data.cpu().numpy())
# print()
lambda_1 = 0.5
lambda_2 = 0.5
loss = first_term + lambda_1 * third_term + lambda_2 * third_term
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
collapsed_softmax_logits = np.argmax(
softmax_shape_net_final_prediction.detach(), axis=1)
collapsed_softmax_unet = np.argmax(
softmax_unet_prediction.detach(), axis=1)
for i in range(1):
rgb_prediction = collapsed_softmax_logits[i].repeat(
3, 1, 1).float()
rgb_prediction = np.moveaxis(rgb_prediction.numpy(), 0, -1)
converted_img = img_to_visible(rgb_prediction)
converted_img = to_tensor(converted_img)
rgb_unet_prediction = collapsed_softmax_unet[i].repeat(
3, 1, 1).float()
rgb_unet_prediction = np.moveaxis(
rgb_unet_prediction.numpy(), 0, -1)
converted_img_unet = img_to_visible(rgb_unet_prediction)
converted_img_unet = to_tensor(converted_img_unet)
masks_changed = masks[i].detach().cpu()
masks_changed = masks_changed.repeat(3, 1, 1).float()
masks_changed = np.moveaxis(masks_changed.numpy(), 0, -1)
masks_changed = img_to_visible(masks_changed)
masks_changed = to_tensor(masks_changed)
# changed_imgs = torch.cat([imgs[i],imgs[i],imgs[i]]).detach().cpu()
# print(changed_imgs.size(), masks_changed.size(), converted_img_unet.size(), converted_img.size())
third_tensor = torch.cat(
(imgs[i].detach().cpu(), masks_changed,
converted_img_unet, converted_img), -1)
if phase == 'val':
writer.add_image(
'ValidationEpoch: {}'.format(str(epoch)),
third_tensor, epoch)
else:
writer.add_image(
'TrainingEpoch: {}'.format(str(epoch)),
third_tensor, epoch)
# statistics
running_loss += loss.detach().item()
batch_IoU = 0.0
for k in range(len(imgs)):
batch_IoU += metrics_evaluator.mean_IU(
collapsed_softmax_logits.numpy()[k],
masks.cpu().numpy()[k])
batch_IoU /= len(imgs)
unet_batch_IoU = 0.0
for j in range(len(imgs)):
unet_batch_IoU += metrics_evaluator.mean_IU(
collapsed_softmax_logits.numpy()[k],
masks.cpu().numpy()[k])
unet_batch_IoU /= len(imgs)
running_IoU += batch_IoU
unet_IoU += unet_batch_IoU
epoch_loss = running_loss / len(data_loader)
epoch_IoU = running_IoU / len(data_loader)
epoch_unet_IoU = unet_IoU / len(data_loader)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss,
epoch_IoU))
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
best_val_loss = epoch_loss
best_IoU = epoch_IoU
best_model_weights = model.state_dict()
if phase == 'val':
# print(optimizer.param_groups[0]['lr'])
curr_val_loss = epoch_loss
curr_val_IoU = epoch_IoU
curr_unet_val_IoU = epoch_unet_IoU
else:
curr_training_loss = epoch_loss
curr_training_IoU = epoch_IoU
curr_unet_training_IoU = epoch_unet_IoU
writer.add_scalars(
'TrainValIoU', {
'trainIoU': curr_training_IoU,
'validationIoU': curr_val_IoU,
'trainUnetIoU': curr_unet_training_IoU,
'validationUnetIoU': curr_unet_val_IoU
}, epoch)
writer.add_scalars('TrainValLoss', {
'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
}, epoch)
# Saving best model
torch.save(
best_model_weights,
os.path.join(
path_to_save_best_weights,
'shape_net_with_pretraining{:2f}.pth'.format(best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
def train(model, train_loader, val_loader, optimizer, num_epochs, path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1)# Use for NLLLoss()
softmax = nn.Softmax(dim=1)
criterion_nlloss = nn.NLLLoss()
criterion_mseloss = nn.MSELoss().to(device)
metrics_evaluator = PerformanceMetricsEvaluator()
to_tensor = transforms.ToTensor()
writer = SummaryWriter('runs/shape_net_with_only_crossentropy/')
since = time.time()
best_model_weights = model.state_dict()
best_IoU = 0.0
best_val_loss = 1000000000
curr_val_loss = 0.0
curr_training_loss = 0.0
curr_training_IoU = 0.0
curr_val_IoU = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step(best_val_loss)
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
running_IoU = 0
# Iterate over data.
ind = 0
for imgs, masks in tqdm(data_loader):
imgs, masks = imgs.to(device), masks.to(device)
masks_for_shape = masks.clone().unsqueeze(1).float()
# zero the parameter gradients
optimizer.zero_grad()
# forward
logits, encoded_shape = model(imgs)
_, encoded_mask = model(masks_for_shape)
log_logits = log_softmax(logits).to(device)
nll_loss = criterion_nlloss(log_logits, masks)
# mse_loss = criterion_mseloss(encoded_shape, encoded_mask)
loss = nll_loss #+ 0.5*mse_loss
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
softmax_logits = softmax(logits)
collapsed_softmax_logits = np.argmax(softmax_logits.detach(), axis=1)
for i in range(len(collapsed_softmax_logits)):
if phase == 'val':
writer.add_image('ValidationEpoch: {}'.format(str(epoch)),
vutils.make_grid([imgs[i].cpu(),masks[i].cpu().float().unsqueeze(0), collapsed_softmax_logits[i].float().unsqueeze(0)]), epoch)
else:
writer.add_image('TrainingEpoch: {}'.format(str(epoch)),
vutils.make_grid([imgs[i].cpu(),masks[i].cpu().float().unsqueeze(0), collapsed_softmax_logits[i].float().unsqueeze(0)]), epoch)
# statistics
running_loss += loss.detach().item()
running_IoU += metrics_evaluator.mean_IU(collapsed_softmax_logits.numpy()[0], masks.cpu().numpy()[0])
ind+=1
epoch_loss = running_loss / len(data_loader)
epoch_IoU = running_IoU / len(data_loader)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_IoU))
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
best_val_loss = epoch_loss
best_IoU = epoch_IoU
best_model_weights = model.state_dict()
if phase == 'val':
# print(optimizer.param_groups[0]['lr'])
curr_val_loss = epoch_loss
curr_val_IoU = epoch_IoU
else:
curr_training_loss = epoch_loss
curr_training_IoU = epoch_IoU
writer.add_scalars('TrainValIoU',
{'trainIoU': curr_training_IoU,
'validationIoU': curr_val_IoU
},
epoch
)
writer.add_scalars('TrainValLoss',
{'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
},
epoch
)
# Saving best model
torch.save(best_model_weights,
os.path.join(path_to_save_best_weights, 'shape_net_with_only_crossentropy{:2f}.pth'.format(best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
def train(model, train_loader, val_loader, optimizer, num_epochs, path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1)# Use for NLLLoss()
softmax = nn.Softmax(dim=1)
# criterion_nlloss = nn.NLLLoss(weight=torch.tensor([0.5015290518452679, 164.00001523734954]).to(device))
metrics_evaluator = PerformanceMetricsEvaluator()
to_tensor = transforms.ToTensor()
writer = SummaryWriter('runs/unet_256/')
since = time.time()
best_model_weights = model.state_dict()
# best_IoU = 0.0
best_dice = 0.0
best_val_loss = 1000000000
curr_val_loss = 0.0
curr_training_loss = 0.0
curr_training_dice = 0.0
curr_val_dice = 0.0
# curr_training_IoU = 0.0
# curr_val_IoU = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step(best_val_loss)
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
# running_IoU = 0
running_dice = 0.
# Iterate over data.
for imgs, masks in tqdm(data_loader):
imgs, masks = imgs.to(device), masks.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
logits = model(imgs)
softmax_logits = softmax(logits)
one_hoted_masks = make_one_hot(torch.unsqueeze(masks, 1))
# loss = soft_dice_coeff(softmax_logits, one_hoted_masks)
loss = 1 - dice_coeff(softmax_logits, one_hoted_masks)
# log_softmax_logits = log_softmax(logits)
# loss = criterion_nlloss(log_softmax_logits, masks)
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
if phase == 'val':
name = 'ValidationEpoch'
else:
name = 'TrainingEpoch'
collapsed_softmax_logits = torch.argmax(softmax_logits, dim=1)
empty_tensor = torch.zeros((collapsed_softmax_logits[0].size())).float()
pred_mask_overlapping = vutils.make_grid([ collapsed_softmax_logits[0].float().cpu(),masks[0].float().cpu(),empty_tensor.cpu()], nrow=3)
image_to_show = vutils.make_grid([imgs[0].cpu(),pred_mask_overlapping])
writer.add_image('{}: {}'.format(name, str(epoch)), image_to_show,epoch)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.detach().item()
# batch_IoU = 0.0
batch_dice = 0.0
for k in range(len(imgs)):
# batch_IoU += metrics_evaluator.mean_IU(collapsed_softmax_logits.numpy()[k], masks.cpu().numpy()[k])
batch_dice += dice_coeff(softmax_logits, one_hoted_masks).item()
# batch_IoU /= len(imgs)
batch_dice /= len(imgs)
running_dice += batch_dice
# running_IoU += batch_IoU
epoch_loss = running_loss / len(data_loader)
# epoch_IoU = running_IoU / len(data_loader)
epoch_dice = running_dice / len(data_loader)
print('{} Loss: {:.4f} Dice: {:.4f}'.format(phase, epoch_loss, epoch_dice))
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
best_val_loss = epoch_loss
# best_IoU = epoch_IoU
best_dice = epoch_dice
best_model_weights = model.state_dict()
# Saving best model
torch.save(best_model_weights, os.path.join(path_to_save_best_weights, 'unet_baseline_dice_256_{:2f}.pth'.format(best_val_loss)))
if phase == 'val':
# print(optimizer.param_groups[0]['lr'])
curr_val_loss = epoch_loss
# curr_val_IoU = epoch_IoU
curr_val_dice = epoch_dice
else:
curr_training_loss = epoch_loss
# curr_training_IoU = epoch_IoU
curr_training_dice = epoch_dice
writer.add_scalars('TrainVal_dice',
{'train_dice': curr_training_dice,
'validation_dice': curr_val_dice
},
epoch
)
writer.add_scalars('TrainValLoss',
{'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
},
epoch
)
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
def train(model, train_loader, val_loader, optimizer, num_epochs,
path_to_save_best_weights):
model.train()
log_softmax = nn.LogSoftmax(dim=1) # Use for NLLLoss()
softmax = nn.Softmax(dim=1)
criterion_nlloss = nn.NLLLoss().to(device)
criterion_mseloss = nn.MSELoss().to(device)
metrics_evaluator = PerformanceMetricsEvaluator()
writer = SummaryWriter('LumenTissueCorruption')
since = time.time()
best_model_weights = model.state_dict()
best_IoU = 0.0 # TODO
curr_val_loss = 0.0
curr_training_loss = 0.0
best_val_loss = 1000000000
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step() TODO the learning rate plato
model.train()
data_loader = train_loader
else:
model.eval()
data_loader = val_loader
running_loss = 0.0
# running_corrects = 0 TODO add IoU
# Iterate over data.
for imgs, masks in tqdm(data_loader):
mask_to_encode = masks.numpy()
mask_to_encode = (
np.arange(4) == mask_to_encode[..., None]).astype(float)
mask_to_encode = np.moveaxis(mask_to_encode, 3, 1)
mask_to_encode = torch.from_numpy(mask_to_encode).float().to(
device)
imgs, masks = imgs.to(device), masks.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
print(imgs.shape)
return
logits, encoded_shape = model(imgs)
_, encoded_mask = model(mask_to_encode)
log_logits = log_softmax(logits)
softmax_logits = softmax(logits).detach().cpu()
nll_loss = criterion_nlloss(log_logits, masks)
# mse_loss1 = criterion_mseloss(softmax_logits, imgs)
mse_loss2 = criterion_mseloss(encoded_shape, encoded_mask)
loss = nll_loss + 0.5 * mse_loss2
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
collapsed_softmax_logits = np.argmax(softmax_logits.detach(),
axis=1)
uncollupsed_images = (np.arange(4) == collapsed_softmax_logits[
..., None]).numpy().astype(np.uint8)
uncollupsed_images = torch.from_numpy(
np.moveaxis(uncollupsed_images, 3, 1)).float()
for i in range(len(softmax_logits)):
empty_channel = np.zeros((544, 544), dtype=np.uint64)
if phase == 'val':
img_name = 'ValidationEpoch: {}'.format(str(epoch))
else:
img_name = 'TrainingEpoch: {}'.format(str(epoch))
writer.add_image(
img_name,
vutils.make_grid([
imgs[i][1:, :, :].detach().cpu(),
softmax_logits[i][1:, :, :],
uncollupsed_images[i][1:, :, :]
]), epoch)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.detach().item()
# running_corrects += TODO
epoch_loss = running_loss / len(data_loader)
# epoch_acc = running_corrects / dataset_sizes[phase] # TODO add IoU
print('{} Loss: {:.4f}'.format(phase, epoch_loss))
# print('{} Loss: {:.4f} Acc: {:.4f}'.format(
# phase, epoch_loss, epoch_acc)) TODO add IoU
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
# best_acc = epoch_acc
best_val_loss = epoch_loss
best_model_weights = model.state_dict()
if phase == 'val':
curr_val_loss = epoch_loss
else:
curr_training_loss = epoch_loss
writer.add_scalars('TrainValLoss', {
'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
}, epoch)
# Saving best model
torch.save(
best_model_weights,
os.path.join(
path_to_save_best_weights,
'Unet_manually_corrupted_lumenTissue{:2f}.pth'.format(
best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
def train(unet, shapeNet, train_loader, val_loader, optimizer1, optimizer2,
num_epochs, path_to_save_best_weights):
unet.train()
shapeNet.train()
log_softmax = nn.LogSoftmax(dim=1) # Use for NLLLoss()
softmax = nn.Softmax(dim=1)
criterion_nlloss = nn.NLLLoss().to(device)
criterion_mseloss1 = nn.MSELoss().to(device)
criterion_mseloss2 = nn.MSELoss().to(device)
metrics_evaluator = PerformanceMetricsEvaluator()
writer = SummaryWriter('FrameworkTensorboard/exp9/')
since = time.time()
best_model_weights = shapeNet.state_dict()
best_IoU = 0.0 # TODO
curr_val_loss = 0.0
curr_training_loss = 0.0
best_val_loss = 1000000000
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
for phase in ['train', 'val']:
if phase == 'train':
# scheduler.step() TODO the learning rate plato
unet.train()
shapeNet.train()
data_loader = train_loader
else:
unet.eval()
shapeNet.eval()
data_loader = val_loader
running_loss = 0.0
# running_corrects = 0 TODO add IoU
# Iterate over data.
for imgs, masks in tqdm(data_loader):
mask_to_encode = masks.numpy()
mask_to_encode = (
np.arange(4) == mask_to_encode[..., None]).astype(float)
mask_to_encode = torch.from_numpy(
np.moveaxis(mask_to_encode, 3, 1)).float().to(device)
imgs, masks = imgs.to(device), masks.to(device)
# zero the parameter gradients
optimizer1.zero_grad()
optimizer2.zero_grad()
# forward
unet_prediction = unet(imgs)
shape_net_final_prediction, shape_net_encoded_prediction = shapeNet(
softmax(unet_prediction))
_, encoded_mask = shapeNet(mask_to_encode)
log_softmax_unet_prediction = log_softmax(unet_prediction)
softmax_unet_prediction = softmax(unet_prediction)
softmax_shape_net_final_prediction = softmax(
shape_net_final_prediction)
first_term = criterion_mseloss1(mask_to_encode,
shape_net_final_prediction)
# second_term = criterion_mseloss(encoded_mask, shape_net_encoded_prediction)
third_term = criterion_mseloss2(unet_prediction,
mask_to_encode)
lambda_1 = 0.5
lambda_2 = 0.5
#loss = third_term + 0.5*first_term
# loss = third_term + 0.2*second_term
# loss = second_term + third_term
# loss = second_term + third_term
# ================================================================== #
# Tensorboard Logging #
# ================================================================== #
unet_softmax_collupsed = np.argmax(
softmax_unet_prediction.detach(), axis=1)
unet_softmax_uncollupsed = (np.arange(
4) == unet_softmax_collupsed[..., None]).numpy().astype(
np.uint8)
unet_softmax_uncollupsed = torch.from_numpy(
np.moveaxis(unet_softmax_uncollupsed, 3, 1)).float()
softmax_shape_net_final_collupsed = np.argmax(
softmax_shape_net_final_prediction.detach(), axis=1)
softmax_shape_net_final_uncollupsed = (
np.arange(4) == softmax_shape_net_final_collupsed[...,
None]
).numpy().astype(np.uint8)
softmax_shape_net_final_uncollupsed = torch.from_numpy(
np.moveaxis(softmax_shape_net_final_uncollupsed, 3,
1)).float()
for i in range(len(unet_softmax_uncollupsed)):
empty_channel = np.zeros((544, 544), dtype=np.uint64)
if phase == 'val':
img_name = 'ValidationEpoch: {}'.format(str(epoch))
else:
img_name = 'TrainingEpoch: {}'.format(str(epoch))
writer.add_image(
img_name,
vutils.make_grid([
imgs[i].detach().cpu(),
unet_softmax_uncollupsed[i][1:, :, :],
softmax_shape_net_final_uncollupsed[i][
1:, :, :].cpu()
]), epoch)
# backward + optimize only if in training phase
if phase == 'train':
# first_term.backward(retain_graph=True)
third_term.backward(retain_graph=True)
first_term.backward()
optimizer1.step()
optimizer2.step()
# statistics
running_loss += first_term.detach().item() + third_term.detach(
).item()
# running_corrects += TODO
epoch_loss = running_loss / len(data_loader)
# epoch_acc = running_corrects / dataset_sizes[phase] # TODO add IoU
print('{} Loss: {:.4f}'.format(phase, epoch_loss))
# print('{} L1: {:.4f}'.format(
# phase, first_term))
# print('{} L2: {:.4f}'.format(
# phase, second_term))
# print('{} L3: {:.4f}'.format(
# phase, third_term))
# print('{} Loss: {:.4f} Acc: {:.4f}'.format(
# phase, epoch_loss, epoch_acc)) TODO add IoU
# deep copy the model
if phase == 'val' and epoch_loss < best_val_loss: # TODO add IoU
# best_acc = epoch_acc
best_val_loss = epoch_loss
best_model_weights = shapeNet.state_dict()
if phase == 'val':
curr_val_loss = epoch_loss
else:
curr_training_loss = epoch_loss
writer.add_scalars('TrainValLoss', {
'trainLoss': curr_training_loss,
'validationLoss': curr_val_loss
}, epoch)
# Saving best model
torch.save(
best_model_weights,
os.path.join(
path_to_save_best_weights,
'Unet_manually_corrupted_lumenTissue{:2f}.pth'.format(
best_val_loss)))
# Show the timing and final statistics
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Loss: {:4f}'.format(best_val_loss)) # TODO add IoU
