-
Notifications
You must be signed in to change notification settings - Fork 0
/
paired_autoencoder_traintest_functions.py
145 lines (121 loc) · 6.69 KB
/
paired_autoencoder_traintest_functions.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/usr/bin/env python3
"""
This file contains functions for paired AutoEncoder model training and testing with PyTorch.
Contents
---
apply_criterion() : calculates latent features and loss given criterion between output and target
calc_fundus_flio_loss() : calculates custom loss weighted by fundus and FLIO losses
train() : trains model
test() : validates model
"""
import numpy as np
import os
import PIL.Image
import torch
from autoencoder_traintest_functions import StopCondition
from loss_utils import cosine_similarity_loss, SSIM, MS_SSIM, MEF_MSSSIM_Loss
from tensorboard_utils import add_image_tensorboard, denormalize, denormalize_and_rescale
def apply_criterion(model, image, target, criterion, save_outputs=None):
"""
:param model: PyTorch model
:param image: input image data
:param target: given input, this is the desired output; clone of image, as reconstruction is the goal
:param criterion: loss type
:param save_outputs: save model output to variable
:return: latent features and loss calculation given criterion between output and target
"""
latent_features, output = model.encoder(image), model(image)
if isinstance(save_outputs, list):
save_outputs.append(output)
if isinstance(criterion, SSIM) or isinstance(criterion, MS_SSIM):
output, target = denormalize(output), denormalize(target)
return latent_features, criterion(output, target)
def calc_fundus_flio_loss(fundus_loss, fundus_latent_features, flio_loss, flio_latent_features, criterion_latent):
"""
:param fundus_loss, flio_loss: loss criterion for each model
:param fundus_latent_features, flio_latent_features: features of latent space for each model
:param criterion_latent: loss criterion for shared constraints and model learning based on paired latent spaces
:return: custom loss weighted by fundus and FLIO losses
"""
return 0.01 * fundus_loss + 0.01 * flio_loss + criterion_latent(fundus_latent_features, flio_latent_features)
def train(model, model2, trainloader, epoch_num, criterion, criterion2, criterion_latent, optimizer, scheduler, device,
writer, output_path, plot_steps=1000, stop_condition=4000):
"""
:param model, model2: paired PyTorch models
:param trainloader: PyTorch DataLoader
:param epoch_num: number of training epochs
:param criterion, criterion2: loss type
:param criterion_latent: loss criterion for shared constraints and model learning based on paired latent spaces
:param optimizer: optimization algorithm for stochastic gradient descent
:param scheduler: learning rate scheduler
:param device: PyTorch device ('cpu' or 'cuda')
:param writer: write to TensorBoard
:param output_path: directory to save figures created during training
:param plot_steps: number of steps to plot training progress in TensorBoard
:param stop_condition: number of steps without improvement for early stopping
"""
model.train()
model2.train()
if stop_condition:
stop_condition = StopCondition(stop_condition)
epoch_log = open("log.txt", "a")
for epoch in range(epoch_num):
running_loss = 0
for i, (fundus_image, flio_image, fundus_target, flio_target, image_files) in enumerate(trainloader):
optimizer.zero_grad()
fundus_image, fundus_target = fundus_image.to(device), fundus_target.to(device)
fundus_latent_features, fundus_loss = apply_criterion(model, fundus_image, fundus_target, criterion)
flio_image, flio_target = flio_image.to(device), flio_target.to(device)
flio_latent_features, flio_loss = apply_criterion(model2, flio_image, flio_target, criterion2)
total_loss = calc_fundus_flio_loss(fundus_loss, fundus_latent_features, flio_loss, flio_latent_features,
criterion_latent)
total_loss.backward()
optimizer.step()
step = i + epoch * len(trainloader)
batch_loss = total_loss.item()
running_loss += batch_loss
# if i % 100 == 0:
if step % plot_steps == 0: # Generate training progress reconstruction figures every # steps
add_image_tensorboard(model, fundus_image, step=step, epoch=epoch, loss=batch_loss, writer=writer,
output_path=output_path)
if i % len(trainloader) == len(trainloader) - 1:
print('[Epoch: {}, i: {}] loss: {:.5f}'.format(epoch + 1, i + 1, running_loss / len(trainloader)))
# TensorBoard
writer.add_scalar('train_scalar', running_loss / len(trainloader), step)
if stop_condition:
stop_condition.evaluate_stop(running_loss)
if stop_condition.stop:
print('Early stop at [Epoch: {}, i: {}] loss: {:.5f}'.format(epoch + 1, i + 1, running_loss / len(trainloader)))
running_loss = 0
# torch.save(model.state_dict(), 'checkpoint.pth')
epoch_log.write('Epoch: ' + str(epoch))
scheduler.step(total_loss)
epoch_log.close()
def test(model, model2, testloader, criterion, criterion2, criterion_latent, device):
"""
:param model, model2: paired PyTorch models
:param testloader: PyTorch DataLoader
:param criterion, criterion2: loss type
:param criterion_latent: loss criterion for shared constraints and model learning based on paired latent spaces
:param device: PyTorch device ('cpu' or 'cuda')
"""
model.eval()
fundus_outputs = []
flio_outputs = []
losses = []
fundus_filenames = []
flio_filenames = []
with torch.no_grad():
for fundus_image, flio_image, fundus_target, flio_target, image_files in testloader:
fundus_image, fundus_target = fundus_image.to(device), fundus_target.to(device)
fundus_latent_features, fundus_loss = apply_criterion(model, fundus_image, fundus_target, criterion,
save_outputs=fundus_outputs)
flio_image, flio_target = flio_image.to(device), flio_target.to(device)
flio_latent_features, flio_loss = apply_criterion(model2, flio_image, flio_target, criterion2,
save_outputs=flio_outputs)
total_loss = calc_fundus_flio_loss(fundus_loss, fundus_latent_features, flio_loss, flio_latent_features,
criterion_latent)
losses.append(total_loss)
fundus_filenames.append(image_files[0])
flio_filenames.append(image_files[1])
return fundus_outputs, flio_outputs, losses, fundus_filenames, flio_filenames