Variational Generator

This commit is contained in:
Si11ium
2020-03-10 16:59:51 +01:00
parent 21e7e31805
commit 1b5a7dc69e
10 changed files with 177 additions and 95 deletions

View File

@ -25,59 +25,33 @@ class CNNRouteGeneratorModel(LightningBaseModule):
return Adam(self.parameters(), lr=self.hparams.train_param.lr)
def training_step(self, batch_xy, batch_nb, *args, **kwargs):
batch_x, label = batch_xy
generated_alternative, z, mu, logvar = self(batch_x + [label, ])
map_array, trajectory = batch_x
map_stack = torch.cat((map_array, trajectory, generated_alternative), dim=1)
pred_label = self.discriminator(map_stack)
discriminated_bce_loss = self.criterion(pred_label, label.float().unsqueeze(-1))
batch_x, alternative = batch_xy
generated_alternative, z, mu, logvar = self(batch_x)
mse_loss = self.criterion(generated_alternative, alternative)
# see Appendix B from VAE paper:
# Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
# https://arxiv.org/abs/1312.6114
# 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
# Dimensional Resizing
# Dimensional Resizing TODO: Does This make sense? Sanity Check it!
kld_loss /= reduce(mul, self.in_shape)
loss = (kld_loss + discriminated_bce_loss) / 2
return dict(loss=loss, log=dict(loss=loss,
discriminated_bce_loss=discriminated_bce_loss,
kld_loss=kld_loss)
)
loss = (kld_loss + mse_loss) / 2
return dict(loss=loss, log=dict(loss=loss, mse_loss=mse_loss, kld_loss=kld_loss))
def _test_val_step(self, batch_xy, batch_nb, *args):
batch_x, label = batch_xy
batch_x, alternative = batch_xy
map_array, trajectory, label = batch_x
generated_alternative, z, mu, logvar = self(batch_x + [label, ])
map_array, trajectory = batch_x
generated_alternative, z, mu, logvar = self(batch_x)
map_stack = torch.cat((map_array, trajectory, generated_alternative), dim=1)
pred_label = self.discriminator(map_stack)
discriminated_bce_loss = self.criterion(pred_label, label.float().unsqueeze(-1))
return dict(discriminated_bce_loss=discriminated_bce_loss, batch_nb=batch_nb,
pred_label=pred_label, label=label, generated_alternative=generated_alternative)
def validation_step(self, *args):
return self._test_val_step(*args)
def validation_epoch_end(self, outputs: list):
return self._test_val_epoch_end(outputs)
return dict(batch_nb=batch_nb, label=label, generated_alternative=generated_alternative, pred_label=-1)
def _test_val_epoch_end(self, outputs, test=False):
evaluation = ROCEvaluation(plot_roc=True)
pred_label = torch.cat([x['pred_label'] for x in outputs])
labels = torch.cat([x['label'] for x in outputs]).unsqueeze(1)
mean_losses = torch.stack([x['discriminated_bce_loss'] for x in outputs]).mean()
# Sci-py call ROC eval call is eval(true_label, prediction)
roc_auc, tpr, fpr = evaluation(labels.cpu().numpy(), pred_label.cpu().numpy(), )
labels = torch.cat([x['label'] for x in outputs]).unsqueeze(1)
if test:
# self.logger.log_metrics(score_dict)
self.logger.log_image(f'{self.name}_ROC-Curve', plt.gcf())
plt.clf()
maps, trajectories, labels, val_restul_dict = self.generate_random()
@ -87,7 +61,13 @@ class CNNRouteGeneratorModel(LightningBaseModule):
fig = g.draw()
self.logger.log_image(f'{self.name}_Output', fig)
return dict(mean_losses=mean_losses, roc_auc=roc_auc, epoch=self.current_epoch)
return dict(epoch=self.current_epoch)
def validation_step(self, *args):
return self._test_val_step(*args)
def validation_epoch_end(self, outputs: list):
return self._test_val_epoch_end(outputs)
def test_step(self, *args):
return self._test_val_step(*args)
@ -95,31 +75,20 @@ class CNNRouteGeneratorModel(LightningBaseModule):
def test_epoch_end(self, outputs):
return self._test_val_epoch_end(outputs, test=True)
@property
def discriminator(self):
if self._disc is None:
raise RuntimeError('Set the Discriminator first; "set_discriminator(disc_model)')
return self._disc
def set_discriminator(self, disc_model):
if self._disc is not None:
raise RuntimeError('Discriminator has already been set... What are trying to do?')
self._disc = disc_model
def __init__(self, *params):
def __init__(self, *params, issubclassed=False):
super(CNNRouteGeneratorModel, self).__init__(*params)
# Dataset
self.dataset = TrajData(self.hparams.data_param.map_root, mode='just_route',
length=self.hparams.data_param.dataset_length)
if not issubclassed:
# Dataset
self.dataset = TrajData(self.hparams.data_param.map_root, mode='separated_arrays',
length=self.hparams.data_param.dataset_length)
self.criterion = nn.MSELoss()
# Additional Attributes
self.in_shape = self.dataset.map_shapes_max
# Todo: Better naming and size in Parameters
self.feature_dim = 10
self.lat_dim = self.feature_dim + self.feature_dim + 1
self._disc = None
# NN Nodes
###################################################
@ -127,7 +96,6 @@ class CNNRouteGeneratorModel(LightningBaseModule):
# Utils
self.relu = nn.ReLU()
self.sigmoid = nn.Sigmoid()
self.criterion = nn.MSELoss()
#
# Map Encoder
@ -222,7 +190,7 @@ class CNNRouteGeneratorModel(LightningBaseModule):
alt_tensor = self.alt_deconv_2(alt_tensor)
alt_tensor = self.alt_deconv_3(alt_tensor)
alt_tensor = self.alt_deconv_out(alt_tensor)
alt_tensor = self.sigmoid(alt_tensor)
# alt_tensor = self.sigmoid(alt_tensor)
return alt_tensor
def encode(self, map_array, trajectory, label):
@ -266,4 +234,100 @@ class CNNRouteGeneratorModel(LightningBaseModule):
maps = self._move_to_model_device(torch.stack(maps))
labels = self._move_to_model_device(torch.as_tensor([0] * n + [1] * n))
return maps, trajectories, labels, self._test_val_step(([maps, trajectories], labels), -9999)
return maps, trajectories, labels, self._test_val_step(((maps, trajectories, labels), None), -9999)
class CNNRouteGeneratorDiscriminated(CNNRouteGeneratorModel):
name = 'CNNRouteGeneratorDiscriminated'
def training_step(self, batch_xy, batch_nb, *args, **kwargs):
batch_x, label = batch_xy
generated_alternative, z, mu, logvar = self(batch_x)
map_array, trajectory = batch_x
map_stack = torch.cat((map_array, trajectory, generated_alternative), dim=1)
pred_label = self.discriminator(map_stack)
discriminated_bce_loss = self.criterion(pred_label, label.float().unsqueeze(-1))
# see Appendix B from VAE paper:
# Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
# https://arxiv.org/abs/1312.6114
# 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
# Dimensional Resizing
kld_loss /= reduce(mul, self.in_shape)
loss = (kld_loss + discriminated_bce_loss) / 2
return dict(loss=loss, log=dict(loss=loss,
discriminated_bce_loss=discriminated_bce_loss,
kld_loss=kld_loss)
)
def _test_val_step(self, batch_xy, batch_nb, *args):
batch_x, label = batch_xy
generated_alternative, z, mu, logvar = self(batch_x)
map_array, trajectory = batch_x
map_stack = torch.cat((map_array, trajectory, generated_alternative), dim=1)
pred_label = self.discriminator(map_stack)
discriminated_bce_loss = self.criterion(pred_label, label.float().unsqueeze(-1))
return dict(discriminated_bce_loss=discriminated_bce_loss, batch_nb=batch_nb,
pred_label=pred_label, label=label, generated_alternative=generated_alternative)
def validation_step(self, *args):
return self._test_val_step(*args)
def validation_epoch_end(self, outputs: list):
return self._test_val_epoch_end(outputs)
def _test_val_epoch_end(self, outputs, test=False):
evaluation = ROCEvaluation(plot_roc=True)
pred_label = torch.cat([x['pred_label'] for x in outputs])
labels = torch.cat([x['label'] for x in outputs]).unsqueeze(1)
mean_losses = torch.stack([x['discriminated_bce_loss'] for x in outputs]).mean()
# Sci-py call ROC eval call is eval(true_label, prediction)
roc_auc, tpr, fpr = evaluation(labels.cpu().numpy(), pred_label.cpu().numpy(), )
if test:
# self.logger.log_metrics(score_dict)
self.logger.log_image(f'{self.name}_ROC-Curve', plt.gcf())
plt.clf()
maps, trajectories, labels, val_restul_dict = self.generate_random()
from lib.visualization.generator_eval import GeneratorVisualizer
g = GeneratorVisualizer(maps, trajectories, labels, val_restul_dict)
fig = g.draw()
self.logger.log_image(f'{self.name}_Output', fig)
return dict(mean_losses=mean_losses, roc_auc=roc_auc, epoch=self.current_epoch)
def test_step(self, *args):
return self._test_val_step(*args)
def test_epoch_end(self, outputs):
return self._test_val_epoch_end(outputs, test=True)
@property
def discriminator(self):
if self._disc is None:
raise RuntimeError('Set the Discriminator first; "set_discriminator(disc_model)')
return self._disc
def set_discriminator(self, disc_model):
if self._disc is not None:
raise RuntimeError('Discriminator has already been set... What are trying to do?')
self._disc = disc_model
def __init__(self, *params):
super(CNNRouteGeneratorDiscriminated, self).__init__(*params, issubclassed=True)
self._disc = None
self.criterion = nn.BCELoss()
self.dataset = TrajData(self.hparams.data_param.map_root, mode='just_route',
length=self.hparams.data_param.dataset_length)

View File

@ -32,24 +32,35 @@ class ConvHomDetector(LightningBaseModule):
pred_y = self(batch_x)
return dict(prediction=pred_y, label=batch_y, batch_nb=batch_nb)
def validation_step(self, batch_xy, batch_nb, **kwargs):
batch_x, batch_y = batch_xy
pred_y = self(batch_x)
return dict(prediction=pred_y, label=batch_y, batch_nb=batch_nb)
def test_epoch_end(self, outputs):
evaluation = ROCEvaluation(plot_roc=True)
return self._val_test_end(outputs)
def validation_epoch_end(self, outputs: list):
return self._val_test_end(outputs)
def _val_test_end(self, outputs, test=True):
evaluation = ROCEvaluation(plot_roc=True if test else False)
predictions = torch.cat([x['prediction'] for x in outputs])
labels = torch.cat([x['label'] for x in outputs]).unsqueeze(1)
# Sci-py call ROC eval call is eval(true_label, prediction)
roc_auc, tpr, fpr = evaluation(labels.cpu().numpy(), predictions.cpu().numpy(), )
score_dict = dict(roc_auc=roc_auc)
roc_auc, tpr, fpr = evaluation(labels.cpu().numpy(), predictions.cpu().numpy())
# self.logger.log_metrics(score_dict)
self.logger.log_image(f'{self.name}', plt.gcf())
if test:
self.logger.log_image(f'{self.name}', plt.gcf())
return dict(log=score_dict)
return dict(score=roc_auc, log=dict(roc_auc=roc_auc))
def __init__(self, hparams):
super(ConvHomDetector, self).__init__(hparams)
# Dataset
self.dataset = TrajData(self.hparams.data_param.map_root, mode='all_in_map')
self.dataset = TrajData(self.hparams.data_param.map_root, mode='all_in_map', )
# Additional Attributes
self.map_shape = self.dataset.map_shapes_max
@ -59,6 +70,7 @@ class ConvHomDetector(LightningBaseModule):
assert len(self.in_shape) == 3, f'Image or map shape has to have 3 dims, but had: {len(self.in_shape)}'
self.criterion = nn.BCELoss()
self.sigmoid = nn.Sigmoid()
self.relu = nn.ReLU()
# NN Nodes
# ============================
@ -100,6 +112,7 @@ class ConvHomDetector(LightningBaseModule):
tensor = self.map_conv_3(tensor)
tensor = self.flatten(tensor)
tensor = self.linear(tensor)
tensor = self.relu(tensor)
tensor = self.classifier(tensor)
tensor = self.sigmoid(tensor)
return tensor

View File

@ -17,7 +17,7 @@ class ConvModule(nn.Module):
output = self(x)
return output.shape[1:]
def __init__(self, in_shape, activation: nn.Module = nn.ELU, pooling_size=None, use_bias=True, use_norm=True,
def __init__(self, in_shape, activation: nn.Module = nn.ELU, pooling_size=None, use_bias=True, use_norm=False,
dropout: Union[int, float] = 0, conv_class=nn.Conv2d,
conv_filters=64, conv_kernel=5, conv_stride=1, conv_padding=0):
super(ConvModule, self).__init__()

View File

@ -154,7 +154,7 @@ class LightningBaseModule(pl.LightningModule, ABC):
# Validation Dataloader
def val_dataloader(self):
return DataLoader(dataset=self.dataset.val_dataset, shuffle=False,
return DataLoader(dataset=self.dataset.val_dataset, shuffle=True,
batch_size=self.hparams.train_param.batch_size,
num_workers=self.hparams.data_param.worker)

View File

@ -18,10 +18,6 @@ import lib.variables as V
class Map(object):
# This setting is for Img mode "L" aka GreyScale Image; values: 0-255
white = 255
black = 0
def __copy__(self):
return copy.deepcopy(self)
@ -51,6 +47,7 @@ class Map(object):
def __init__(self, name='', array_like_map_representation=None):
if array_like_map_representation is not None:
array_like_map_representation = array_like_map_representation.astype(np.float32)
if array_like_map_representation.ndim == 2:
array_like_map_representation = np.expand_dims(array_like_map_representation, axis=0)
assert array_like_map_representation.ndim == 3
@ -70,7 +67,7 @@ class Map(object):
# Check pixels for their color (determine if walkable)
for idx, value in np.ndenumerate(self.map_array):
if value != self.black:
if value != V.BLACK:
# IF walkable, add node
graph.add_node(idx, count=0)
# Fully connect to all surrounding neighbors
@ -91,10 +88,9 @@ class Map(object):
if image.mode != 'L':
image = image.convert('L')
map_array = np.expand_dims(np.array(image), axis=0)
map_array = np.where(np.asarray(map_array) == cls.white, 1, 0)
if embedding_size:
assert isinstance(embedding_size, tuple), f'embedding_size was of type: {type(embedding_size)}'
embedding = np.zeros(embedding_size)
embedding = np.full(embedding_size, V.BLACK)
embedding[:map_array.shape[0], :map_array.shape[1], :map_array.shape[2]] = map_array
map_array = embedding
@ -146,12 +142,15 @@ class Map(object):
polyline = trajectory.xy_vertices
polyline.extend(reversed(other_trajectory.xy_vertices))
img = Image.new('L', (self.height, self.width), 0)
img = Image.new('L', (self.height, self.width), color=V.WHITE)
draw = ImageDraw.Draw(img)
draw.polygon(polyline, outline=1, fill=1)
draw.polygon(polyline, outline=V.BLACK, fill=V.BLACK)
a = (np.asarray(img) * np.where(self.as_2d_array == self.black, 1, 0)).sum()
binary_img = np.where(np.asarray(img).squeeze() == V.BLACK, 1, 0)
binary_map = np.where(self.as_2d_array == V.BLACK, 1, 0)
a = (binary_img * binary_map).sum()
if a:
return V.ALTERNATIVE # Non-Homotoph

View File

@ -42,9 +42,9 @@ class Trajectory(object):
return list(zip(self._vertices[:-1], self._vertices[1:]))
def draw_in_array(self, shape):
trajectory_space = np.zeros(shape)
trajectory_space = np.zeros(shape).astype(np.float32)
for index in self.vertices:
trajectory_space[index] = 1
trajectory_space[index] = V.WHITE
return trajectory_space
@property

View File

@ -5,7 +5,7 @@ from collections import defaultdict
from configparser import ConfigParser
from pathlib import Path
from lib.models.generators.cnn import CNNRouteGeneratorModel
from lib.models.generators.cnn import CNNRouteGeneratorModel, CNNRouteGeneratorDiscriminated
from lib.models.homotopy_classification.cnn_based import ConvHomDetector
from lib.utils.model_io import ModelParameters
@ -28,7 +28,10 @@ class Config(ConfigParser):
@property
def model_class(self):
model_dict = dict(classifier_cnn=ConvHomDetector, generator_cnn=CNNRouteGeneratorModel)
model_dict = dict(ConvHomDetector=ConvHomDetector,
CNNRouteGenerator=CNNRouteGeneratorModel,
CNNRouteGeneratorDiscriminated=CNNRouteGeneratorDiscriminated
)
try:
return model_dict[self.get('model', 'type')]
except KeyError as e:

View File

@ -3,3 +3,5 @@ _ROOT = Path('..')
HOMOTOPIC = 1
ALTERNATIVE = 0
WHITE = 255
BLACK = 0