|
import functools |
|
import os |
|
from collections.abc import Iterable |
|
from datetime import datetime |
|
|
|
import jax |
|
import jax.numpy as jnp |
|
import numpy as np |
|
import optax |
|
import pandas as pd |
|
import tqdm |
|
from flax.training import train_state |
|
from tensorboardX import SummaryWriter |
|
|
|
from jax_nca.utils import make_circle_masks |
|
|
|
|
|
def get_tensorboard_logger( |
|
experiment_name: str, base_log_path: str = "tensorboard_logs" |
|
): |
|
log_path = "{}/{}_{}".format(base_log_path, experiment_name, datetime.now()) |
|
train_writer = SummaryWriter(log_path, flush_secs=10) |
|
full_log_path = os.path.join(os.getcwd(), log_path) |
|
print( |
|
"Follow tensorboard logs with: python -m tensorboard.main --logdir '{}'".format( |
|
full_log_path |
|
) |
|
) |
|
return train_writer |
|
|
|
|
|
def create_train_state(rng, nca, learning_rate, shape): |
|
nca_seed = nca.create_seed( |
|
nca.num_hidden_channels, nca.num_target_channels, shape=shape[:-1], batch_size=1 |
|
) |
|
"""Creates initial `TrainState`.""" |
|
params = nca.init(rng, nca_seed, rng)["params"] |
|
tx = optax.chain( |
|
|
|
optax.adam(learning_rate), |
|
) |
|
return train_state.TrainState.create(apply_fn=nca.apply, params=params, tx=tx) |
|
|
|
|
|
def clip_grad_norm(grad): |
|
factor = 1.0 / ( |
|
jnp.linalg.norm(jax.tree_util.tree_leaves(jax.tree_map(jnp.linalg.norm, grad))) |
|
+ 1e-8 |
|
) |
|
return jax.tree_map((lambda x: x * factor), grad) |
|
|
|
|
|
@functools.partial(jax.jit, static_argnames=("apply_fn", "num_steps")) |
|
def train_step( |
|
apply_fn, state, seeds: jnp.array, targets: jnp.array, num_steps: int, rng |
|
): |
|
def mse_loss(pred, y): |
|
squared_diff = jnp.square(pred - y) |
|
return jnp.mean(squared_diff, axis=[-3, -2, -1]) |
|
|
|
def loss_fn(params): |
|
def forward(carry, inp): |
|
carry = apply_fn({"params": params}, carry, rng) |
|
return carry, carry |
|
|
|
x, outs = jax.lax.scan(forward, seeds, None, length=num_steps) |
|
rgb, a = x[..., :3], jnp.clip(x[..., 3:4], 0.0, 1.0) |
|
rgb = jnp.clip(1.0 - a + rgb, 0.0, 1.0) |
|
|
|
outs = jnp.transpose(outs, [1, 0, 2, 3, 4]) |
|
subset = outs[:, -8:] |
|
return jnp.mean( |
|
jax.vmap(mse_loss)(subset[..., :4], jnp.expand_dims(targets, 1)) |
|
), (x, rgb) |
|
|
|
grad_fn = jax.value_and_grad(loss_fn, has_aux=True) |
|
(loss, aux), grads = grad_fn(state.params) |
|
grads = clip_grad_norm(grads) |
|
updated, rgb = aux |
|
return state.apply_gradients(grads=grads), loss, grads, updated, rgb |
|
|
|
|
|
class SamplePool: |
|
def __init__(self, max_size: int = 1000): |
|
self.max_size = max_size |
|
self.pool = [None] * max_size |
|
|
|
def __getitem__(self, idx): |
|
if isinstance(idx, Iterable): |
|
return [self.pool[i] for i in idx] |
|
return idx |
|
|
|
def __setitem__(self, idx, v): |
|
if isinstance(idx, Iterable): |
|
for i in range(len(idx)): |
|
index = idx[i] |
|
self.pool[index] = v[i] |
|
else: |
|
self.pool[idx] = v |
|
|
|
def sample(self, num_samples: int): |
|
indices = np.random.randint(0, self.max_size, num_samples) |
|
return self.__getitem__(indices), indices |
|
|
|
|
|
def flatten(d): |
|
df = pd.json_normalize(d, sep="_") |
|
return df.to_dict(orient="records")[0] |
|
|
|
|
|
class EmojiTrainer: |
|
def __init__(self, dataset, nca, pool_size: int = 1024, n_damage: int = 0): |
|
self.dataset = dataset |
|
self.img_shape = self.dataset.img_shape |
|
self.nca = nca |
|
self.pool_size = pool_size |
|
self.n_damage = n_damage |
|
self.state = None |
|
|
|
def train( |
|
self, |
|
num_epochs, |
|
batch_size: int = 8, |
|
seed: int = 10, |
|
lr: float = 0.001, |
|
min_steps: int = 64, |
|
max_steps: int = 96, |
|
): |
|
pool = SamplePool(self.pool_size) |
|
|
|
writer = get_tensorboard_logger("EMOJITrainer") |
|
rng = jax.random.PRNGKey(seed) |
|
rng, init_rng = jax.random.split(rng) |
|
self.state = create_train_state(init_rng, self.nca, lr, self.dataset.img_shape) |
|
|
|
bar = tqdm.tqdm(np.arange(num_epochs)) |
|
try: |
|
for i in bar: |
|
num_steps = int(np.random.randint(min_steps, max_steps)) |
|
samples, indices = pool.sample(batch_size) |
|
for j in range(len(samples)): |
|
if samples[j] is None: |
|
samples[j] = self.nca.create_seed( |
|
self.nca.num_hidden_channels, |
|
self.nca.num_target_channels, |
|
shape=self.img_shape[:-1], |
|
batch_size=1, |
|
)[0] |
|
samples[0] = self.nca.create_seed( |
|
self.nca.num_hidden_channels, |
|
self.nca.num_target_channels, |
|
shape=self.img_shape[:-1], |
|
batch_size=1, |
|
)[0] |
|
batch = np.stack(samples) |
|
if self.n_damage > 0: |
|
damage = ( |
|
1.0 |
|
- make_circle_masks( |
|
int(self.n_damage), self.img_shape[0], self.img_shape[1] |
|
)[..., None] |
|
) |
|
batch[-self.n_damage :] *= damage |
|
|
|
batch = jnp.array(batch) |
|
targets, rgb_targets = self.dataset.get_batch(batch_size) |
|
targets = jnp.array(targets) |
|
|
|
self.state, loss, grads, outputs, rgb_outputs = train_step( |
|
self.nca.apply, |
|
self.state, |
|
batch, |
|
targets, |
|
num_steps=num_steps, |
|
rng=rng, |
|
) |
|
|
|
grad_dict = {k: dict(grads[k]) for k in grads.keys()} |
|
grad_dict = flatten(grad_dict) |
|
|
|
grad_dict = { |
|
k: {kk: np.sum(vv).item() for kk, vv in v.items()} |
|
for k, v in grad_dict.items() |
|
} |
|
grad_dict = flatten(grad_dict) |
|
|
|
pool[indices] = np.array(outputs) |
|
|
|
bar.set_description("Loss: {}".format(loss.item())) |
|
|
|
self.emit_metrics( |
|
writer, |
|
i, |
|
batch, |
|
rgb_outputs, |
|
rgb_targets, |
|
loss.item(), |
|
metrics=grad_dict, |
|
) |
|
|
|
return self.state |
|
except Exception: |
|
return self.state |
|
|
|
def emit_metrics( |
|
self, train_writer, i: int, batch, outputs, targets, loss, metrics={} |
|
): |
|
train_writer.add_scalar("loss", loss, i) |
|
|
|
train_writer.add_images("batch", self.nca.to_rgb(batch), i, dataformats="NHWC") |
|
train_writer.add_images("outputs", outputs, i, dataformats="NHWC") |
|
train_writer.add_images("targets", targets, i, dataformats="NHWC") |
|
for k in metrics: |
|
train_writer.add_scalar(k, metrics[k], i) |
|
|
