-
Notifications
You must be signed in to change notification settings - Fork 11
/
task.py
258 lines (201 loc) · 9.34 KB
/
task.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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
import torch
import torch.nn as nn
import torch.utils.data as data
import numpy as np
import math
from random import shuffle, random
from statistics import mean, median
import matplotlib.pyplot as plt
import os
import sys
import time
from statistics import stdev
sys.path.append("../")
from general.mutils import debug_level, append_in_dict, get_device
class TaskTemplate:
def __init__(self, model, model_params, name, load_data=True, debug=False, batch_size=64, drop_last=False, num_workers=None):
# Saving parameters
self.name = name
self.model = model
self.model_params = model_params
self.batch_size = batch_size
self.train_batch_size = batch_size
self.debug = debug
# Initializing dataset parameters
self.train_dataset = None
self.val_dataset = None
self.test_dataset = None
self.train_epoch = 0
# Load data if specified, and create data loaders
if load_data:
self._load_datasets()
self._initialize_data_loaders(drop_last=drop_last, num_workers=num_workers)
else:
self.train_data_loader = None
self.train_data_loader_iter = None
self.val_data_loader = None
self.test_data_loader = None
# Create a dictionary to store summary metrics in
self.summary_dict = {}
# Placeholders for visualization
self.gen_batch = None
self.class_colors = None
# Put model on correct device
self.model.to(get_device())
def _initialize_data_loaders(self, drop_last, num_workers):
if num_workers is None:
if isinstance(self.model, nn.DataParallel) and torch.cuda.device_count() > 1:
num_workers = torch.cuda.device_count()
else:
num_workers = 1
def _init_fn(worker_id):
np.random.seed(42)
# num_workers = 1
# Initializes all data loaders with the loaded datasets
if hasattr(self.train_dataset, "get_sampler"):
self.train_data_loader = data.DataLoader(self.train_dataset, batch_sampler=self.train_dataset.get_sampler(self.train_batch_size, drop_last=drop_last), pin_memory=True,
num_workers=num_workers, worker_init_fn=_init_fn)
self.val_data_loader = data.DataLoader(self.val_dataset, batch_sampler=self.val_dataset.get_sampler(self.train_batch_size, drop_last=False), pin_memory=True, num_workers=1, worker_init_fn=_init_fn)
self.test_data_loader = data.DataLoader(self.test_dataset, batch_sampler=self.test_dataset.get_sampler(self.train_batch_size, drop_last=False), pin_memory=True, num_workers=1, worker_init_fn=_init_fn)
else:
self.train_data_loader = data.DataLoader(self.train_dataset, batch_size=self.train_batch_size, shuffle=True, pin_memory=True, drop_last=drop_last, num_workers=num_workers,
worker_init_fn=_init_fn)
self.val_data_loader = data.DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, pin_memory=True, drop_last=False, num_workers=1, worker_init_fn=_init_fn)
self.test_data_loader = data.DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False, pin_memory=True, drop_last=False, num_workers=1, worker_init_fn=_init_fn)
self.train_data_loader_iter = iter(self.train_data_loader) # Needed to retrieve batch by batch from dataset
def train_step(self, iteration=0):
# Check if training data was correctly loaded
if self.train_data_loader_iter is None:
print("[!] ERROR: Iterator of the training data loader was None. Additional parameters: " + \
"train_data_loader was %sloaded, " % ("not " if self.train_data_loader is None else "") + \
"train_dataset was %sloaded." % ("not " if self.train_dataset is None else ""))
# Get batch and put it on correct device
batch = self._get_next_batch()
batch = TaskTemplate.batch_to_device(batch)
# Perform task-specific training step
return self._train_batch(batch, iteration=iteration)
def eval(self, data_loader=None, **kwargs):
# Default: if no dataset is specified, we use validation dataset
if data_loader is None:
assert self.val_data_loader is not None, "[!] ERROR: Validation dataset not loaded. Please load the dataset beforehand for evaluation."
data_loader = self.val_data_loader
is_test = (data_loader == self.test_data_loader)
start_time = time.time()
torch.cuda.empty_cache()
self.model.eval()
# Prepare metrics
total_nll, embed_nll, nll_counter = 0, 0, 0
# Evaluation loop
with torch.no_grad():
for batch_ind, batch in enumerate(data_loader):
if debug_level() == 0:
print("Evaluation process: %4.2f%%" % (100.0 * batch_ind / len(data_loader)), end="\r")
# Put batch on correct device
batch = TaskTemplate.batch_to_device(batch)
# Evaluate single batch
batch_size = batch[0].size(0) if isinstance(batch, tuple) else batch.size(0)
batch_nll = self._eval_batch(batch, is_test=is_test)
total_nll += batch_nll.item() * batch_size
nll_counter += batch_size
if self.debug and batch_ind > 10:
break
avg_nll = total_nll / max(1e-5, nll_counter)
detailed_metrics = {
"negative_log_likelihood": avg_nll,
"bpd": self.loss_to_bpd(avg_nll), # Bits per dimension
}
with torch.no_grad():
self._eval_finalize_metrics(detailed_metrics, is_test=is_test, **kwargs)
self.model.train()
eval_time = int(time.time() - start_time)
print("Finished %s with bpd of %4.3f (%imin %is)" % ("testing" if data_loader == self.test_data_loader else "evaluation", detailed_metrics["bpd"], eval_time/60, eval_time%60))
torch.cuda.empty_cache()
if "loss_metric" in detailed_metrics:
loss_metric = detailed_metrics["loss_metric"]
else:
loss_metric = avg_nll
return loss_metric, detailed_metrics
def loss_to_bpd(self, loss):
return (np.log2(np.exp(1)) * loss)
def test(self, **kwargs):
return self.eval(data_loader=self.test_data_loader, **kwargs)
def add_summary(self, writer, iteration, checkpoint_path=None):
# Adding metrics collected during training to the tensorboard
# Function can/should be extended if needed
for key, val in self.summary_dict.items():
summary_key = "train_%s/%s" % (self.name, key)
if not isinstance(val, list): # If it is not a list, it is assumably a single scalar
writer.add_scalar(summary_key, val, iteration)
self.summary_dict[key] = 0.0
elif len(val) == 0: # Skip an empty list
continue
elif not isinstance(val[0], list): # For a list of scalars, report the mean
writer.add_scalar(summary_key, mean(val), iteration)
self.summary_dict[key] = list()
else: # List of lists indicates a histogram
val = [v for sublist in val for v in sublist]
writer.add_histogram(summary_key, np.array(val), iteration)
self.summary_dict[key] = list()
def _ldj_per_layer_to_summary(self, ldj_per_layer, pre_phrase="ldj_layer_"):
for layer_index, layer_ldj in enumerate(ldj_per_layer):
if isinstance(layer_ldj, tuple) or isinstance(layer_ldj, list):
for i in range(len(layer_ldj)):
append_in_dict(self.summary_dict, "ldj_layer_%i_%i" % (layer_index, i), layer_ldj[i].detach().mean().item())
elif isinstance(layer_ldj, dict):
for key, ldj_val in layer_ldj.items():
append_in_dict(self.summary_dict, "ldj_layer_%i_%s" % (layer_index, key), ldj_val.detach().mean().item())
else:
append_in_dict(self.summary_dict, "ldj_layer_%i" % layer_index, layer_ldj.detach().mean().item())
def _get_next_batch(self):
# Try to get next batch. If one epoch is over, the iterator throws an error, and we start a new iterator
try:
batch = next(self.train_data_loader_iter)
except StopIteration:
self.train_data_loader_iter = iter(self.train_data_loader)
batch = next(self.train_data_loader_iter)
self.train_epoch += 1
return batch
#######################################################
### Abstract method to be implemented by subclasses ###
#######################################################
def _load_datasets(self):
# Function for initializing datasets. Should set the following class parameters:
# -> self.train_dataset
# -> self.val_dataset
# -> self.test_dataset
raise NotImplementedError
def _train_batch(self, batch, iteration=0):
# Given a batch, return the loss to be trained on
raise NotImplementedError
def _eval_batch(self, batch, is_test=False, take_mean=True):
# Given a batch, return the negative log likelihood for its elements
# Input arguments:
# -> "is_test": True during testing, False during validation
# -> "take_mean": If true, the return should be the average log
# likelihood of the batch. Otherwise, return per element.
raise NotImplementedError
def finalize_summary(self, writer, iteration, checkpoint_path):
# This function is called after the finishing training and performing testing.
# Can be used to add something to the summary before finishing.
pass
def export_best_results(self, checkpoint_path, iteration):
# This function is called if the last evaluation has been the best so far.
# Can be used to add some output to the checkpoint directory.
pass
def initialize(self):
# This function is called before starting the training.
pass
def _eval_finalize_metrics(self, detailed_metrics, is_test=False, initial_eval=False):
# This function is called after evaluation finished.
# Can be used to add final metrics to the dictionary, which is added to the tensorboard.
pass
####################
## Static methods ##
####################
@staticmethod
def batch_to_device(batch):
if isinstance(batch, tuple) or isinstance(batch, list):
batch = tuple([b.to(get_device()) for b in batch])
else:
batch = batch.to(get_device())
return batch