def test_real_histo_data(tmpdir):
logger = Logger(str(tmpdir), flush_secs=0.1)
logger.log_histogram('hist2', [1, 7, 6, 9, 8, 1, 4, 5, 3, 7], step=1)
logger.log_histogram('hist2', [5, 3, 2, 0, 8, 5, 7, 7, 7, 2], step=2)
logger.log_histogram('hist2', [1, 2, 2, 1, 5, 1, 8, 4, 4, 1], step=3)
tf_log, = glob.glob(str(tmpdir) + '/*')
assert os.path.basename(tf_log).startswith('events.out.tfevents.')
def test_dummy_histo():
logger = Logger(None, is_dummy=True)
bins = [0, 1, 2, 3]
logger.log_histogram('key', (bins, [0.0, 1.0, 2.0]), step=1)
logger.log_histogram('key', (bins, [1.0, 1.5, 2.5]), step=2)
logger.log_histogram('key', (bins, [0.0, 1.0, 2.0]), step=3)
assert dict(logger.dummy_log) == {
'key': [(1, (bins, [0.0, 1.0, 2.0])), (2, (bins, [1.0, 1.5, 2.5])),
(3, (bins, [0.0, 1.0, 2.0]))]
}
def test_real_histo_tuple(tmpdir):
"""
from tests.test_tensorboard_logger import *
import ubelt as ub
ub.delete(ub.ensure_app_cache_dir('tf_logger'))
tmpdir = ub.ensure_app_cache_dir('tf_logger/runs/run1')
"""
logger = Logger(str(tmpdir), flush_secs=0.1)
bins = [-.5, .5, 1.5, 2.5]
logger.log_histogram('hist1', (bins, [0.0, 1.0, 2.0]), step=1)
logger.log_histogram('hist1', (bins, [1.0, 1.5, 2.5]), step=2)
logger.log_histogram('hist1', (bins, [0.0, 1.0, 2.0]), step=3)
tf_log, = glob.glob(str(tmpdir) + '/*')
assert os.path.basename(tf_log).startswith('events.out.tfevents.')
def _save(self, checkpoint_dir):
file_path = checkpoint_dir + '/ray_SNN_W_epoch_' + str(
self.epoch) + '.pickle'
# self.model.save_weights(file_path)
if self.epoch % 20 == 0:
with open(file_path, 'w') as fw:
save = {'w_h': self.model.w_h, 'w_o': self.model.w_o}
pickle.dump(save, fw)
# writer = SummaryWriter(checkpoint_dir)
logger = Logger(checkpoint_dir + '/images/')
# logger_hist = Logger(checkpoint_dir + '/histograms/')
# img = np.random.rand(10, 10)
images = []
# mx = np.max(self.w_h[0])
# mn = np.min(self.w_h[0])
for tp in range(self.model.outputs):
tpp = self.model.w_o[:, tp]
sz = np.ceil((np.sqrt(np.size(tpp)))).astype(int)
tpm = np.zeros(sz * sz, dtype=float)
tpm[0:len(tpp)] = tpp
tpp = tpm
tpp = np.reshape(tpp, newshape=[sz, sz])
images.append(scale(tpp))
cimages = combine_matrix(*images)
logger.log_images('key_out', [cimages], step=self.epoch)
logger.log_histogram('key_out', [self.model.w_o], step=self.epoch)
for k in range(len(self.h)):
images = []
for tp in range(self.h[k]):
tpp = self.model.w_h[k][:, tp]
sz = np.ceil((np.sqrt(np.size(tpp)))).astype(int)
tpm = np.zeros(sz * sz, dtype=float)
tpm[0:np.size(tpp)] = tpp
tpp = tpm
tpp = np.reshape(tpp, newshape=[sz, sz])
images.append(scale(tpp))
cimages = combine_matrix(*images)
logger.log_images('key_' + str(k), [cimages], step=self.epoch)
logger.log_histogram('key' + str(k), [self.model.w_h[k]],
step=self.epoch)
return file_path
class TrainNetwork:
def __init__(self,
batch_size=64,
image_size=64,
load_model=None,
iterations_start=0,
seed=None,
epochs=10000,
lr_decay_iter=250000,
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=1e-3,
weighting_factor=0.5,
regression=False,
loss='MultiLabelSoftMarginLoss',
network='ColorfulImageColorization',
convert_on_gpu=False,
do_not_log=False):
# hyperparameters
self.batch_size = batch_size
self.image_size = image_size
self.iterations = iterations_start
self.seed = np.random.randint(0, 10000) if seed is None else seed
self.epochs = epochs
self.lr_decay_iter = lr_decay_iter
self.lr0 = lr
self.betas = betas
self.eps = eps
self.weight_decay = weight_decay
self.weighting_factor = weighting_factor
self.lr = self.calc_learning_rate()
self.regression = regression
self.loss = loss
self.model_name = '{date:%Y_%m_%d__%H_%M_%S}_{net}_{loss}_{dataset}'.format(
date=datetime.datetime.now(),
net=network,
loss=self.loss,
dataset=dataset)
torch.manual_seed(self.seed)
self.convert_on_gpu = convert_on_gpu
# tensorboard
if not do_not_log:
self.logger = Logger(logs_path + self.model_name)
self.log_hyperparameter()
# model, loss, optimizer
assert network in [
'DeepKoalarization', 'CheapConvNet', 'ColorfulImageColorization',
'DeepKoalarizationNorm'
]
if self.regression:
out_channels = 2
else:
out_channels = int((256 / grid_size)**2)
if network == 'DeepKoalarization':
self.model = DeepKoalarization(
out_channels=out_channels,
to_rgb=(self.loss == 'PerceptualLoss'))
elif network == 'DeepKoalarizationNorm':
self.model = DeepKoalarizationNorm(
out_channels=out_channels,
to_rgb=(self.loss == 'PerceptualLoss'))
elif network == 'CheapConvNet':
assert self.loss != 'PerceptualLoss'
self.model = CheapConvNet(out_channels=out_channels)
elif network == 'ColorfulImageColorization':
self.model = ColorfulImageColorization(
out_channels=out_channels,
to_rgb=(self.loss == 'PerceptualLoss'))
self.optimizer = optim.Adam(self.model.parameters(),
lr=self.lr,
betas=self.betas,
eps=self.eps,
weight_decay=self.weight_decay)
assert loss in [
'PerceptualLoss', 'MSELoss', 'MultiLabelSoftMarginLoss',
'BCEWithLogitsLoss', 'MultinomialCrossEntropyLoss'
]
if loss == 'PerceptualLoss':
self.loss_fn = PerceptualLoss()
elif loss == 'MSELoss':
assert self.regression
self.loss_fn = nn.MSELoss()
elif loss == 'MultiLabelSoftMarginLoss':
assert not self.regression
w = torch.load(images_path +
'classification_weights_{}_{}.pth'.format(
grid_size, self.weighting_factor))
self.loss_fn = nn.MultiLabelSoftMarginLoss(weight=w.view(-1, 1, 1))
elif loss == 'BCEWithLogitsLoss':
assert not self.regression
w = torch.load(images_path +
'classification_weights_{}_{}.pth'.format(
grid_size, self.weighting_factor))
self.loss_fn = nn.BCEWithLogitsLoss(weight=w.view(-1, 1, 1))
elif loss == 'MultinomialCrossEntropyLoss':
assert not self.regression
w = torch.load(images_path +
'classification_weights_{}_{}.pth'.format(
grid_size, self.weighting_factor))
self.loss_fn = MultinomialCrossEntropyLoss(weights=w)
# load model
if load_model is not None:
# Load pre learned AlexNet with changed number of output classes
state_dict = torch.load(trained_models_path + load_model,
map_location='cpu')
self.model.load_state_dict(state_dict['model'])
self.optimizer.load_state_dict(state_dict['optimizer'])
self.adjust_learning_rate()
# Use cuda if available
self.cuda = torch.cuda.is_available()
if self.cuda:
self.model.cuda()
self.loss_fn.cuda()
if load_model is not None:
for state in self.optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
# Load dataset
kwargs = {'num_workers': 8, 'pin_memory': True} if self.cuda else {}
self.train_loader = torch.utils.data.DataLoader(ColorizationDataset(
images_path,
train=True,
size=(self.image_size, self.image_size),
target_rgb=(loss == 'PerceptualLoss'),
convert_to_categorical=(not self.regression),
do_not_convert=convert_on_gpu),
batch_size=batch_size,
shuffle=True,
**kwargs)
kwargs = {'num_workers': 1, 'pin_memory': True} if self.cuda else {}
self.test_loader = torch.utils.data.DataLoader(ColorizationDataset(
images_path,
train=False,
size=(self.image_size, self.image_size),
target_rgb=(loss == 'PerceptualLoss'),
convert_to_categorical=(not self.regression),
do_not_convert=convert_on_gpu),
batch_size=8,
drop_last=True,
shuffle=True,
**kwargs)
self.test_iterator = iter(self.test_loader)
def calc_learning_rate(self):
"""
Reduce the learning rate by factor 0.5 every lr_decay_iter
:return: None
"""
lr = self.lr0 * (0.1**(self.iterations // self.lr_decay_iter))
return lr
def adjust_learning_rate(self):
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.lr
def reduce_learning_rate(self):
lr = self.calc_learning_rate()
if abs(lr - self.lr) > 1e-7:
self.lr = lr
self.adjust_learning_rate()
def train(self):
"""
Train the model for one epoch and save the result as a .pth file
:return: None
"""
self.model.train()
for epoch in range(1, self.epochs + 1):
train_loss_epoche = 0
batch_start_time = time.clock()
batch_idx = 0
for batch_idx, (data, target) in enumerate(self.train_loader):
self.reduce_learning_rate()
train_loss_epoche += self.train_one_iter(
data, target, epoch, batch_idx)
self.iterations += 1
print('Batch ' + str(batch_idx + 1) + ' took ' +
str(time.clock() - batch_start_time) + ' seconds')
batch_start_time = time.clock()
# Print information about current epoch
train_loss_epoche /= (batch_idx + 1)
print('Train Epoch: {} \tAverage loss: {:.6f}'.format(
epoch, train_loss_epoche))
def train_one_iter(self, data, target, epoch, batch_idx):
if self.cuda:
data = data.cuda()
target = target.cuda()
if self.convert_on_gpu:
lab = pdc.rgb2lab(data.float() / 255)
data, target = torch.split(lab, [1, 2], dim=1)
if not self.regression:
target = conversion_batch(target)
# Optimize using backpropagation
self.optimizer.zero_grad()
output = self.model(data)
loss = self.loss_fn(output, target)
loss.backward()
self.optimizer.step()
# Print information about current step
print('Train Epoch: {} [{}/{}]\tLoss: {:.6f}'.format(
epoch, batch_idx + 1, len(self.train_loader), loss.item()))
if self.iterations % 5 == 0:
# log loss
test_data, test_target = self.get_next_test_batch()
self.model.eval()
test_output = self.model(test_data)
self.model.train()
test_loss = self.loss_fn(test_output, test_target)
self.log_scalars(self.iterations, loss, test_loss)
if self.iterations % 50 == 0:
# log images
self.log_images(self.iterations, test_data, test_target,
test_output, data, target, output)
if self.iterations % 1000 == 0 and self.iterations > 0:
# Save snapshot
model_name = self.model_name + '_iter{}'.format(self.iterations)
torch.save(
{
'model': self.model.state_dict(),
'optimizer': self.optimizer.state_dict()
}, trained_models_path + '{}.pth'.format(model_name))
if self.iterations % 600 == 0:
# log stuff
self.log_values_gradients(self.iterations)
return loss.item()
def log_hyperparameter(self):
info = {
'batch_size': self.batch_size,
'image_size': self.image_size,
'seed': self.seed,
'epochs': self.epochs,
'learning_decay_iter': self.lr_decay_iter,
'learning_rate_0': self.lr0,
'betas[0]': self.betas[0],
'betas[1]': self.betas[1],
'eps_optimizer': self.eps,
'weight_decay_optimizer': self.weight_decay,
'weighting_factor': self.weighting_factor,
'regression': self.regression,
'convert_on_gpu': self.convert_on_gpu
}
for tag, value in info.items():
self.logger.log_value(tag, value, 0)
def log_scalars(self, step, train_loss, test_loss):
# adapted from https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/04-utils/tensorboard/main.py
# 1. Log scalar values (scalar summary)
info = {
'train_loss': train_loss.item(),
'test_loss': test_loss.item(),
'learning_rate': self.lr
}
for tag, value in info.items():
self.logger.log_value(tag, value, step)
def log_images(self, step, test_data, test_target, test_output, train_data,
train_target, train_output):
# 3. Log test images (image summary)
num_images = 1
test_original = self.convert_to_images(test_data[:num_images],
test_target[:num_images],
is_target=True)
test_colorized = self.convert_to_images(test_data[:num_images],
test_output[:num_images],
is_target=False)
train_original = self.convert_to_images(train_data[:num_images],
train_target[:num_images],
is_target=True)
train_colorized = self.convert_to_images(train_data[:num_images],
train_output[:num_images],
is_target=False)
info = {
'test colorized': test_colorized,
'test original': test_original,
'train colorized': train_colorized,
'train original': train_original
}
for tag, images in info.items():
self.logger.log_images(tag, images, step)
def log_values_gradients(self, step):
# adapted from https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/04-utils/tensorboard/main.py
# 2. Log values and gradients of the parameters (histogram summary)
for tag, value in self.model.named_parameters():
if 'feature_extractor' not in tag:
tag = tag.replace('.', '/')
self.logger.log_histogram(tag, value.data.cpu().numpy(), step)
self.logger.log_histogram(tag + '/grad',
value.grad.data.cpu().numpy(), step)
def convert_to_images(self, data, output_or_target, is_target, t=0.38):
if self.loss == 'PerceptualLoss':
return output_or_target.detach().cpu().permute(
0, 2, 3, 1).numpy().astype(np.float64)
else:
return generate_images_numpy(data,
output_or_target,
is_target,
regression=self.regression,
t=t)
def get_next_test_batch(self):
try:
data, target = next(self.test_iterator)
except StopIteration:
self.test_iterator = iter(self.test_loader)
data, target = next(self.test_iterator)
if self.cuda:
data = data.cuda()
target = target.cuda()
if self.convert_on_gpu:
lab = pdc.rgb2lab(data.float() / 255)
data, target = torch.split(lab, [1, 2], dim=1)
if not self.regression:
target = conversion_batch(target)
return data, target
