Upload 14 files

data_frame.py

@@ -0,0 +1,104 @@
+import copy
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+
+class DataFrame(object):
+
+    """Minimal pd.DataFrame analog for handling n-dimensional numpy matrices with additional
+    support for shuffling, batching, and train/test splitting.
+
+    Args:
+        columns: List of names corresponding to the matrices in data.
+        data: List of n-dimensional data matrices ordered in correspondence with columns.
+            All matrices must have the same leading dimension.  Data can also be fed a list of
+            instances of np.memmap, in which case RAM usage can be limited to the size of a
+            single batch.
+    """
+
+    def __init__(self, columns, data):
+        assert len(columns) == len(data), 'columns length does not match data length'
+
+        lengths = [mat.shape[0] for mat in data]
+        assert len(set(lengths)) == 1, 'all matrices in data must have same first dimension'
+
+        self.length = lengths[0]
+        self.columns = columns
+        self.data = data
+        self.dict = dict(zip(self.columns, self.data))
+        self.idx = np.arange(self.length)
+
+    def shapes(self):
+        return pd.Series(dict(zip(self.columns, [mat.shape for mat in self.data])))
+
+    def dtypes(self):
+        return pd.Series(dict(zip(self.columns, [mat.dtype for mat in self.data])))
+
+    def shuffle(self):
+        np.random.shuffle(self.idx)
+
+    def train_test_split(self, train_size, random_state=np.random.randint(1000), stratify=None):
+        train_idx, test_idx = train_test_split(
+            self.idx,
+            train_size=train_size,
+            random_state=random_state,
+            stratify=stratify
+        )
+        train_df = DataFrame(copy.copy(self.columns), [mat[train_idx] for mat in self.data])
+        test_df = DataFrame(copy.copy(self.columns), [mat[test_idx] for mat in self.data])
+        return train_df, test_df
+
+    def batch_generator(self, batch_size, shuffle=True, num_epochs=10000, allow_smaller_final_batch=False):
+        epoch_num = 0
+        while epoch_num < num_epochs:
+            if shuffle:
+                self.shuffle()
+
+            for i in range(0, self.length + 1, batch_size):
+                batch_idx = self.idx[i: i + batch_size]
+                if not allow_smaller_final_batch and len(batch_idx) != batch_size:
+                    break
+                yield DataFrame(
+                    columns=copy.copy(self.columns),
+                    data=[mat[batch_idx].copy() for mat in self.data]
+                )
+
+            epoch_num += 1
+
+    def iterrows(self):
+        for i in self.idx:
+            yield self[i]
+
+    def mask(self, mask):
+        return DataFrame(copy.copy(self.columns), [mat[mask] for mat in self.data])
+
+    def concat(self, other_df):
+        mats = []
+        for column in self.columns:
+            mats.append(np.concatenate([self[column], other_df[column]], axis=0))
+        return DataFrame(copy.copy(self.columns), mats)
+
+    def items(self):
+        return self.dict.items()
+
+    def __iter__(self):
+        return self.dict.items().__iter__()
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, key):
+        if isinstance(key, str):
+            return self.dict[key]
+
+        elif isinstance(key, int):
+            return pd.Series(dict(zip(self.columns, [mat[self.idx[key]] for mat in self.data])))
+
+    def __setitem__(self, key, value):
+        assert value.shape[0] == len(self), 'matrix first dimension does not match'
+        if key not in self.columns:
+            self.columns.append(key)
+            self.data.append(value)
+        self.dict[key] = value

demo.py

@@ -0,0 +1,62 @@
+import numpy as np
+from hand import Hand
+
+import lyrics
+
+
+if __name__ == '__main__':
+    hand = Hand()
+
+    # usage demo
+    lines = [
+      "what's up how are you marcos ?"
+    ]
+    biases = [.75 for i in lines]
+    styles = [9 for i in lines]
+    stroke_colors = ['red', 'green', 'black', 'blue']
+    stroke_widths = [1, 2, 1, 2]
+
+    hand.write(
+        filename='img/usage_demo.svg',
+        lines=lines,
+        biases=biases,
+        styles=styles,
+        stroke_colors=stroke_colors,
+        stroke_widths=stroke_widths
+    )
+
+    # demo number 1 - fixed bias, fixed style
+    lines = lyrics.all_star.split("\n")
+    biases = [.75 for i in lines]
+    styles = [12 for i in lines]
+
+    hand.write(
+        filename='img/all_star.svg',
+        lines=lines,
+        biases=biases,
+        styles=styles,
+    )
+
+    # demo number 2 - fixed bias, varying style
+    lines = lyrics.downtown.split("\n")
+    biases = [.75 for i in lines]
+    styles = np.cumsum(np.array([len(i) for i in lines]) == 0).astype(int)
+
+    hand.write(
+        filename='img/downtown.svg',
+        lines=lines,
+        biases=biases,
+        styles=styles,
+    )
+
+    # demo number 3 - varying bias, fixed style
+    lines = lyrics.give_up.split("\n")
+    biases = .2*np.flip(np.cumsum([len(i) == 0 for i in lines]), 0)
+    styles = [7 for i in lines]
+
+    hand.write(
+        filename='img/give_up.svg',
+        lines=lines,
+        biases=biases,
+        styles=styles,
+    )

drawing.py

@@ -0,0 +1,216 @@
+from __future__ import print_function
+from collections import defaultdict
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.signal import savgol_filter
+from scipy.interpolate import interp1d
+
+
+alphabet = [
+    '\x00', ' ', '!', '"', '#', "'", '(', ')', ',', '-', '.',
+    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';',
+    '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K',
+    'L', 'M', 'N', 'O', 'P', 'R', 'S', 'T', 'U', 'V', 'W', 'Y',
+    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l',
+    'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x',
+    'y', 'z'
+]
+alphabet_ord = list(map(ord, alphabet))
+alpha_to_num = defaultdict(int, list(map(reversed, enumerate(alphabet))))
+num_to_alpha = dict(enumerate(alphabet_ord))
+
+MAX_STROKE_LEN = 1200
+MAX_CHAR_LEN = 75
+
+
+def align(coords):
+    """
+    corrects for global slant/offset in handwriting strokes
+    """
+    coords = np.copy(coords)
+    X, Y = coords[:, 0].reshape(-1, 1), coords[:, 1].reshape(-1, 1)
+    X = np.concatenate([np.ones([X.shape[0], 1]), X], axis=1)
+    offset, slope = np.linalg.inv(X.T.dot(X)).dot(X.T).dot(Y).squeeze()
+    theta = np.arctan(slope)
+    rotation_matrix = np.array(
+        [[np.cos(theta), -np.sin(theta)],
+         [np.sin(theta), np.cos(theta)]]
+    )
+    coords[:, :2] = np.dot(coords[:, :2], rotation_matrix) - offset
+    return coords
+
+
+def skew(coords, degrees):
+    """
+    skews strokes by given degrees
+    """
+    coords = np.copy(coords)
+    theta = degrees * np.pi/180
+    A = np.array([[np.cos(-theta), 0], [np.sin(-theta), 1]])
+    coords[:, :2] = np.dot(coords[:, :2], A)
+    return coords
+
+
+def stretch(coords, x_factor, y_factor):
+    """
+    stretches strokes along x and y axis
+    """
+    coords = np.copy(coords)
+    coords[:, :2] *= np.array([x_factor, y_factor])
+    return coords
+
+
+def add_noise(coords, scale):
+    """
+    adds gaussian noise to strokes
+    """
+    coords = np.copy(coords)
+    coords[1:, :2] += np.random.normal(loc=0.0, scale=scale, size=coords[1:, :2].shape)
+    return coords
+
+
+def encode_ascii(ascii_string):
+    """
+    encodes ascii string to array of ints
+    """
+    return np.array(list(map(lambda x: alpha_to_num[x], ascii_string)) + [0])
+
+
+def denoise(coords):
+    """
+    smoothing filter to mitigate some artifacts of the data collection
+    """
+    coords = np.split(coords, np.where(coords[:, 2] == 1)[0] + 1, axis=0)
+    new_coords = []
+    for stroke in coords:
+        if len(stroke) != 0:
+            x_new = savgol_filter(stroke[:, 0], 7, 3, mode='nearest')
+            y_new = savgol_filter(stroke[:, 1], 7, 3, mode='nearest')
+            xy_coords = np.hstack([x_new.reshape(-1, 1), y_new.reshape(-1, 1)])
+            stroke = np.concatenate([xy_coords, stroke[:, 2].reshape(-1, 1)], axis=1)
+            new_coords.append(stroke)
+
+    coords = np.vstack(new_coords)
+    return coords
+
+
+def interpolate(coords, factor=2):
+    """
+    interpolates strokes using cubic spline
+    """
+    coords = np.split(coords, np.where(coords[:, 2] == 1)[0] + 1, axis=0)
+    new_coords = []
+    for stroke in coords:
+
+        if len(stroke) == 0:
+            continue
+
+        xy_coords = stroke[:, :2]
+
+        if len(stroke) > 3:
+            f_x = interp1d(np.arange(len(stroke)), stroke[:, 0], kind='cubic')
+            f_y = interp1d(np.arange(len(stroke)), stroke[:, 1], kind='cubic')
+
+            xx = np.linspace(0, len(stroke) - 1, factor*(len(stroke)))
+            yy = np.linspace(0, len(stroke) - 1, factor*(len(stroke)))
+
+            x_new = f_x(xx)
+            y_new = f_y(yy)
+
+            xy_coords = np.hstack([x_new.reshape(-1, 1), y_new.reshape(-1, 1)])
+
+        stroke_eos = np.zeros([len(xy_coords), 1])
+        stroke_eos[-1] = 1.0
+        stroke = np.concatenate([xy_coords, stroke_eos], axis=1)
+        new_coords.append(stroke)
+
+    coords = np.vstack(new_coords)
+    return coords
+
+
+def normalize(offsets):
+    """
+    normalizes strokes to median unit norm
+    """
+    offsets = np.copy(offsets)
+    offsets[:, :2] /= np.median(np.linalg.norm(offsets[:, :2], axis=1))
+    return offsets
+
+
+def coords_to_offsets(coords):
+    """
+    convert from coordinates to offsets
+    """
+    offsets = np.concatenate([coords[1:, :2] - coords[:-1, :2], coords[1:, 2:3]], axis=1)
+    offsets = np.concatenate([np.array([[0, 0, 1]]), offsets], axis=0)
+    return offsets
+
+
+def offsets_to_coords(offsets):
+    """
+    convert from offsets to coordinates
+    """
+    return np.concatenate([np.cumsum(offsets[:, :2], axis=0), offsets[:, 2:3]], axis=1)
+
+
+def draw(
+        offsets,
+        ascii_seq=None,
+        align_strokes=True,
+        denoise_strokes=True,
+        interpolation_factor=None,
+        save_file=None
+):
+    strokes = offsets_to_coords(offsets)
+
+    if denoise_strokes:
+        strokes = denoise(strokes)
+
+    if interpolation_factor is not None:
+        strokes = interpolate(strokes, factor=interpolation_factor)
+
+    if align_strokes:
+        strokes[:, :2] = align(strokes[:, :2])
+
+    fig, ax = plt.subplots(figsize=(12, 3))
+
+    stroke = []
+    for x, y, eos in strokes:
+        stroke.append((x, y))
+        if eos == 1:
+            coords = zip(*stroke)
+            ax.plot(coords[0], coords[1], 'k')
+            stroke = []
+    if stroke:
+        coords = zip(*stroke)
+        ax.plot(coords[0], coords[1], 'k')
+        stroke = []
+
+    ax.set_xlim(-50, 600)
+    ax.set_ylim(-40, 40)
+
+    ax.set_aspect('equal')
+    plt.tick_params(
+        axis='both',
+        left='off',
+        top='off',
+        right='off',
+        bottom='off',
+        labelleft='off',
+        labeltop='off',
+        labelright='off',
+        labelbottom='off'
+    )
+
+    if ascii_seq is not None:
+        if not isinstance(ascii_seq, str):
+            ascii_seq = ''.join(list(map(chr, ascii_seq)))
+        plt.title(ascii_seq)
+
+    if save_file is not None:
+        plt.savefig(save_file)
+        print('saved to {}'.format(save_file))
+    else:
+        plt.show()
+    plt.close('all')

hand.py

@@ -0,0 +1,150 @@
+import drawing
+from rnn import rnn
+
+
+import numpy as np
+import svgwrite
+
+
+import logging
+import os
+
+
+class Hand(object):
+
+    def __init__(self):
+        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+        self.nn = rnn(
+            log_dir='logs',
+            checkpoint_dir='checkpoints',
+            prediction_dir='predictions',
+            learning_rates=[.0001, .00005, .00002],
+            batch_sizes=[32, 64, 64],
+            patiences=[1500, 1000, 500],
+            beta1_decays=[.9, .9, .9],
+            validation_batch_size=32,
+            optimizer='rms',
+            num_training_steps=100000,
+            warm_start_init_step=17900,
+            regularization_constant=0.0,
+            keep_prob=1.0,
+            enable_parameter_averaging=False,
+            min_steps_to_checkpoint=2000,
+            log_interval=20,
+            logging_level=logging.CRITICAL,
+            grad_clip=10,
+            lstm_size=400,
+            output_mixture_components=20,
+            attention_mixture_components=10
+        )
+        self.nn.restore()
+
+    def write(self, filename, lines, biases=None, styles=None, stroke_colors=None, stroke_widths=None):
+        valid_char_set = set(drawing.alphabet)
+        for line_num, line in enumerate(lines):
+            if len(line) > 75:
+                raise ValueError(
+                    (
+                        "Each line must be at most 75 characters. "
+                        "Line {} contains {}"
+                    ).format(line_num, len(line))
+                )
+
+            for char in line:
+                if char not in valid_char_set:
+                    raise ValueError(
+                        (
+                            "Invalid character {} detected in line {}. "
+                            "Valid character set is {}"
+                        ).format(char, line_num, valid_char_set)
+                    )
+
+        strokes = self._sample(lines, biases=biases, styles=styles)
+        self._draw(strokes, lines, filename, stroke_colors=stroke_colors, stroke_widths=stroke_widths)
+
+    def _sample(self, lines, biases=None, styles=None):
+        num_samples = len(lines)
+        max_tsteps = 40*max([len(i) for i in lines])
+        biases = biases if biases is not None else [0.5]*num_samples
+
+        x_prime = np.zeros([num_samples, 1200, 3])
+        x_prime_len = np.zeros([num_samples])
+        chars = np.zeros([num_samples, 120])
+        chars_len = np.zeros([num_samples])
+
+        if styles is not None:
+            for i, (cs, style) in enumerate(zip(lines, styles)):
+                x_p = np.load('styles/style-{}-strokes.npy'.format(style))
+                c_p = np.load('styles/style-{}-chars.npy'.format(style)).tostring().decode('utf-8')
+
+                c_p = str(c_p) + " " + cs
+                c_p = drawing.encode_ascii(c_p)
+                c_p = np.array(c_p)
+
+                x_prime[i, :len(x_p), :] = x_p
+                x_prime_len[i] = len(x_p)
+                chars[i, :len(c_p)] = c_p
+                chars_len[i] = len(c_p)
+
+        else:
+            for i in range(num_samples):
+                encoded = drawing.encode_ascii(lines[i])
+                chars[i, :len(encoded)] = encoded
+                chars_len[i] = len(encoded)
+
+        [samples] = self.nn.session.run(
+            [self.nn.sampled_sequence],
+            feed_dict={
+                self.nn.prime: styles is not None,
+                self.nn.x_prime: x_prime,
+                self.nn.x_prime_len: x_prime_len,
+                self.nn.num_samples: num_samples,
+                self.nn.sample_tsteps: max_tsteps,
+                self.nn.c: chars,
+                self.nn.c_len: chars_len,
+                self.nn.bias: biases
+            }
+        )
+        samples = [sample[~np.all(sample == 0.0, axis=1)] for sample in samples]
+        return samples
+
+    def _draw(self, strokes, lines, filename, stroke_colors=None, stroke_widths=None):
+        stroke_colors = stroke_colors or ['black']*len(lines)
+        stroke_widths = stroke_widths or [2]*len(lines)
+
+        line_height = 60
+        view_width = 1000
+        view_height = line_height*(len(strokes) + 1)
+
+        dwg = svgwrite.Drawing(filename=filename)
+        dwg.viewbox(width=view_width, height=view_height)
+        dwg.add(dwg.rect(insert=(0, 0), size=(view_width, view_height), fill='white'))
+
+        initial_coord = np.array([0, -(3*line_height / 4)])
+        for offsets, line, color, width in zip(strokes, lines, stroke_colors, stroke_widths):
+
+            if not line:
+                initial_coord[1] -= line_height
+                continue
+
+            offsets[:, :2] *= 1.5
+            strokes = drawing.offsets_to_coords(offsets)
+            strokes = drawing.denoise(strokes)
+            strokes[:, :2] = drawing.align(strokes[:, :2])
+
+            strokes[:, 1] *= -1
+            strokes[:, :2] -= strokes[:, :2].min() + initial_coord
+            strokes[:, 0] += (view_width - strokes[:, 0].max()) / 2
+
+            prev_eos = 1.0
+            p = "M{},{} ".format(0, 0)
+            for x, y, eos in zip(*strokes.T):
+                p += '{}{},{} '.format('M' if prev_eos == 1.0 else 'L', x, y)
+                prev_eos = eos
+            path = svgwrite.path.Path(p)
+            path = path.stroke(color=color, width=width, linecap='round').fill("none")
+            dwg.add(path)
+
+            initial_coord[1] -= line_height
+
+        dwg.save()

handwriting_api.py

@@ -0,0 +1,52 @@
+from fastapi import FastAPI as FastAPIApp, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import uvicorn
+import os
+from hand import Hand
+from typing import List, Optional
+
+app = FastAPIApp()
+hand = Hand()
+
+class InputData(BaseModel):
+    text: str
+    bias: Optional[float] = 0.75
+    style: Optional[int] = 9
+    stroke_colors: Optional[List[str]] = ['black']
+    stroke_widths: Optional[List[float]] = [2]
+
+def validate_input(data: InputData):
+    if len(data.text) > 75:
+        raise ValueError("Text must be 75 characters or less")
+    if not (0.5 <= data.bias <= 1.0):
+        raise ValueError("Bias must be between 0.5 and 1.0")
+    if not (0 <= data.style <= 12):
+        raise ValueError("Style must be between 0 and 12")
+    if len(data.stroke_colors) != len(data.text.split('\n')):
+        raise ValueError("Number of stroke colors must match number of lines")
+    if len(data.stroke_widths) != len(data.text.split('\n')):
+        raise ValueError("Number of stroke widths must match number of lines")
+
+@app.post("/synthesize")
+def synthesize(data: InputData):
+    try:
+        validate_input(data)
+        lines = data.text.split('\n')
+        biases = [data.bias] * len(lines)
+        styles = [data.style] * len(lines)
+        
+        hand.write(
+            filename='img/output.svg',
+            lines=lines,
+            biases=biases,
+            styles=styles,
+            stroke_colors=data.stroke_colors,
+            stroke_widths=data.stroke_widths
+        )
+        return {"result": "Handwriting synthesized successfully", "output_file": "img/output.svg"}
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

lyrics.py

@@ -0,0 +1,190 @@
+"""lyrics taken from https://www.azlyrics.com/"""
+
+all_star = """Somebody once told me the world is gonna roll me
+I ain't the sharpest tool in the shed
+She was looking kind of dumb with her finger and her thumb
+In the shape of an "L" on her forehead
+
+Well, the years start coming and they don't stop coming
+Fed to the rules and I hit the ground running
+Didn't make sense not to live for fun
+Your brain gets smart but your head gets dumb
+
+So much to do, so much to see
+So what's wrong with taking the back streets?
+You'll never know if you don't go
+You'll never shine if you don't glow
+
+Hey, now, you're an All Star, get your game on, go play
+Hey, now, you're a Rock Star, get the show on, get paid
+And all that glitters is gold
+Only shooting stars break the mold
+
+It's a cool place and they say it gets colder
+You're bundled up now wait 'til you get older
+But the meteor men beg to differ
+Judging by the hole in the satellite picture
+
+The ice we skate is getting pretty thin
+The water's getting warm so you might as well swim
+My world's on fire. How about yours?
+That's the way I like it and I'll never get bored.
+
+Somebody once asked could I spare some change for gas
+I need to get myself away from this place
+I said yep, what a concept
+I could use a little fuel myself
+And we could all use a little change
+
+Well, the years start coming and they don't stop coming
+Fed to the rules and I hit the ground running
+Didn't make sense not to live for fun
+Your brain gets smart but your head gets dumb
+
+So much to do, so much to see
+So what's wrong with taking the back streets?
+You'll never know if you don't go
+You'll never shine if you don't glow.
+
+And all that glitters is gold
+Only shooting stars break the mold"""
+
+downtown = """Making my way downtown
+Walking fast
+Faces pass
+And I'm home-bound
+
+Staring blankly ahead
+Just making my way
+Making a way
+Through the crowd
+
+And I need you
+And I miss you
+And now I wonder
+
+If I could fall into the sky
+Do you think time would pass me by?
+'Cause you know I'd walk a thousand miles
+If I could just see you tonight
+
+It's always times like these
+When I think of you
+And I wonder if you ever think of me
+'Cause everything's so wrong
+And I don't belong
+Living in your precious memory
+
+'Cause I need you
+And I miss you
+And now I wonder
+
+If I could fall into the sky
+Do you think time would pass me by?
+'Cause you know I'd walk a thousand miles
+If I could just see you tonight
+
+And I, I don't wanna let you know
+I, I drown in your memory
+I, I don't wanna let this go
+I, I don't
+
+Making my way downtown
+Walking fast
+Faces pass
+And I'm home-bound
+
+Staring blankly ahead
+Just making my way
+Making a way
+Through the crowd
+
+And I still need you
+And I still miss you
+And now I wonder
+
+If I could fall into the sky
+Do you think time would pass us by?
+'Cause you know I'd walk a thousand miles
+If I could just see you
+
+If I could fall into the sky
+Do you think time would pass me by?
+'Cause you know I'd walk a thousand miles
+If I could just see you
+If I could just hold you tonight"""
+
+give_up = """We're no strangers to love
+You know the rules and so do I
+A full commitment's what I'm thinking of
+You wouldn't get this from any other guy
+
+I just wanna tell you how I'm feeling
+Gotta make you understand
+
+Never gonna give you up
+Never gonna let you down
+Never gonna run around and desert you
+Never gonna make you cry
+Never gonna say goodbye
+Never gonna tell a lie and hurt you
+
+We've known each other for so long
+Your heart's been aching, but
+You're too shy to say it
+Inside, we both know what's been going on
+We know the game and we're gonna play it
+
+And if you ask me how I'm feeling
+Don't tell me you're too blind to see
+
+Never gonna give you up
+Never gonna let you down
+Never gonna run around and desert you
+Never gonna make you cry
+Never gonna say goodbye
+Never gonna tell a lie and hurt you
+
+Never gonna give you up
+Never gonna let you down
+Never gonna run around and desert you
+Never gonna make you cry
+Never gonna say goodbye
+Never gonna tell a lie and hurt you
+
+(Ooh, give you up)
+(Ooh, give you up)
+Never gonna give, never gonna give
+(Give you up)
+Never gonna give, never gonna give
+(Give you up)
+
+We've known each other for so long
+Your heart's been aching, but
+You're too shy to say it
+Inside, we both know what's been going on
+We know the game and we're gonna play it
+
+I just wanna tell you how I'm feeling
+Gotta make you understand
+
+Never gonna give you up
+Never gonna let you down
+Never gonna run around and desert you
+Never gonna make you cry
+Never gonna say goodbye
+Never gonna tell a lie and hurt you
+
+Never gonna give you up
+Never gonna let you down
+Never gonna run around and desert you
+Never gonna make you cry
+Never gonna say goodbye
+Never gonna tell a lie and hurt you
+
+Never gonna give you up
+Never gonna let you down
+Never gonna run around and desert you
+Never gonna make you cry
+Never gonna say goodbye
+Never gonna tell a lie and hurt you"""

prepare_data.py

@@ -0,0 +1,134 @@
+from __future__ import print_function
+import os
+from xml.etree import ElementTree
+
+import numpy as np
+
+import drawing
+
+
+def get_stroke_sequence(filename):
+    tree = ElementTree.parse(filename).getroot()
+    strokes = [i for i in tree if i.tag == 'StrokeSet'][0]
+
+    coords = []
+    for stroke in strokes:
+        for i, point in enumerate(stroke):
+            coords.append([
+                int(point.attrib['x']),
+                -1*int(point.attrib['y']),
+                int(i == len(stroke) - 1)
+            ])
+    coords = np.array(coords)
+
+    coords = drawing.align(coords)
+    coords = drawing.denoise(coords)
+    offsets = drawing.coords_to_offsets(coords)
+    offsets = offsets[:drawing.MAX_STROKE_LEN]
+    offsets = drawing.normalize(offsets)
+    return offsets
+
+
+def get_ascii_sequences(filename):
+    sequences = open(filename, 'r').read()
+    sequences = sequences.replace(r'%%%%%%%%%%%', '\n')
+    sequences = [i.strip() for i in sequences.split('\n')]
+    lines = sequences[sequences.index('CSR:') + 2:]
+    lines = [line.strip() for line in lines if line.strip()]
+    lines = [drawing.encode_ascii(line)[:drawing.MAX_CHAR_LEN] for line in lines]
+    return lines
+
+
+def collect_data():
+    fnames = []
+    for dirpath, dirnames, filenames in os.walk('data/raw/ascii/'):
+        if dirnames:
+            continue
+        for filename in filenames:
+            if filename.startswith('.'):
+                continue
+            fnames.append(os.path.join(dirpath, filename))
+
+    # low quality samples (selected by collecting samples to
+    # which the trained model assigned very low likelihood)
+    blacklist = set(np.load('data/blacklist.npy'))
+
+    stroke_fnames, transcriptions, writer_ids = [], [], []
+    for i, fname in enumerate(fnames):
+        print(i, fname)
+        if fname == 'data/raw/ascii/z01/z01-000/z01-000z.txt':
+            continue
+
+        head, tail = os.path.split(fname)
+        last_letter = os.path.splitext(fname)[0][-1]
+        last_letter = last_letter if last_letter.isalpha() else ''
+
+        line_stroke_dir = head.replace('ascii', 'lineStrokes')
+        line_stroke_fname_prefix = os.path.split(head)[-1] + last_letter + '-'
+
+        if not os.path.isdir(line_stroke_dir):
+            continue
+        line_stroke_fnames = sorted([f for f in os.listdir(line_stroke_dir)
+                                     if f.startswith(line_stroke_fname_prefix)])
+        if not line_stroke_fnames:
+            continue
+
+        original_dir = head.replace('ascii', 'original')
+        original_xml = os.path.join(original_dir, 'strokes' + last_letter + '.xml')
+        tree = ElementTree.parse(original_xml)
+        root = tree.getroot()
+
+        general = root.find('General')
+        if general is not None:
+            writer_id = int(general[0].attrib.get('writerID', '0'))
+        else:
+            writer_id = int('0')
+
+        ascii_sequences = get_ascii_sequences(fname)
+        assert len(ascii_sequences) == len(line_stroke_fnames)
+
+        for ascii_seq, line_stroke_fname in zip(ascii_sequences, line_stroke_fnames):
+            if line_stroke_fname in blacklist:
+                continue
+
+            stroke_fnames.append(os.path.join(line_stroke_dir, line_stroke_fname))
+            transcriptions.append(ascii_seq)
+            writer_ids.append(writer_id)
+
+    return stroke_fnames, transcriptions, writer_ids
+
+
+if __name__ == '__main__':
+    print('traversing data directory...')
+    stroke_fnames, transcriptions, writer_ids = collect_data()
+
+    print('dumping to numpy arrays...')
+    x = np.zeros([len(stroke_fnames), drawing.MAX_STROKE_LEN, 3], dtype=np.float32)
+    x_len = np.zeros([len(stroke_fnames)], dtype=np.int16)
+    c = np.zeros([len(stroke_fnames), drawing.MAX_CHAR_LEN], dtype=np.int8)
+    c_len = np.zeros([len(stroke_fnames)], dtype=np.int8)
+    w_id = np.zeros([len(stroke_fnames)], dtype=np.int16)
+    valid_mask = np.zeros([len(stroke_fnames)], dtype=np.bool)
+
+    for i, (stroke_fname, c_i, w_id_i) in enumerate(zip(stroke_fnames, transcriptions, writer_ids)):
+        if i % 200 == 0:
+            print(i, '\t', '/', len(stroke_fnames))
+        x_i = get_stroke_sequence(stroke_fname)
+        valid_mask[i] = ~np.any(np.linalg.norm(x_i[:, :2], axis=1) > 60)
+
+        x[i, :len(x_i), :] = x_i
+        x_len[i] = len(x_i)
+
+        c[i, :len(c_i)] = c_i
+        c_len[i] = len(c_i)
+
+        w_id[i] = w_id_i
+
+    if not os.path.isdir('data/processed'):
+        os.makedirs('data/processed')
+
+    np.save('data/processed/x.npy', x[valid_mask])
+    np.save('data/processed/x_len.npy', x_len[valid_mask])
+    np.save('data/processed/c.npy', c[valid_mask])
+    np.save('data/processed/c_len.npy', c_len[valid_mask])
+    np.save('data/processed/w_id.npy', w_id[valid_mask])

readme.md

@@ -0,0 +1,68 @@
+![](img/banner.svg)
+# Handwriting Synthesis
+Implementation of the handwriting synthesis experiments in the paper <a href="https://arxiv.org/abs/1308.0850">Generating Sequences with Recurrent Neural Networks</a> by Alex Graves.  The implementation closely follows the original paper, with a few slight deviations, and the generated samples are of similar quality to those presented in the paper.
+
+Web demo from original github (which I forked) is available <a href="https://seanvasquez.com/handwriting-generation/">here</a>.
+
+## Fork Notes
+Fork has the following changes from the original
+1. Updated to work with TensorFlow 2.15.0 (current as of December 2023), but uses the v1 compat mechanism.  Honestly I hacked and burned through many errors in an hour and a half (I forked at 7:05 PM and am writing this at 8:35 PM) and just verified by running the demo.py and looking through the output image files.  It didn't create a banner.svg file (no idea if it was supposed to), and it gets a lot of deprecation warnings so use at your own risk.  TensorFlow will likely abandon some of this V1 compat stuff in the relatively near future, but this should work as long as you can still get 2.15.0 for whatever python version you have.
+2. I'm going to split the Hand class into its own file (per the original author's suggestion).
+3. I left everything else below this text alone (except striking the "split Hand class" request).
+4. If you think my fork is super sloppy (you're right) and want to do it right- the major job is to convert the deprecated `tf.nn.rnn_cell.LSTMCell` and replace it with `tf.keras.layers.LSTMCell`.  It is not a drop in replacement.  Have at it.
+
+## Usage
+```python
+lines = [
+    "Now this is a story all about how",
+    "My life got flipped turned upside down",
+    "And I'd like to take a minute, just sit right there",
+    "I'll tell you how I became the prince of a town called Bel-Air",
+]
+biases = [.75 for i in lines]
+styles = [9 for i in lines]
+stroke_colors = ['red', 'green', 'black', 'blue']
+stroke_widths = [1, 2, 1, 2]
+
+hand = Hand()
+hand.write(
+    filename='img/usage_demo.svg',
+    lines=lines,
+    biases=biases,
+    styles=styles,
+    stroke_colors=stroke_colors,
+    stroke_widths=stroke_widths
+)
+```
+![](img/usage_demo.svg)
+
+~~Currently, the `Hand` class must be imported from `demo.py`.  If someone would like to package this project to make it more usable, please [contribute](#contribute).~~
+
+A pretrained model is included, but if you'd like to train your own, read <a href='https://github.com/sjvasquez/handwriting-synthesis/tree/master/data/raw'>these instructions</a>.
+
+## Demonstrations
+Below are a few hundred samples from the model, including some samples demonstrating the effect of priming and biasing the model.  Loosely speaking, biasing controls the neatness of the samples and priming controls the style of the samples. The code for these demonstrations can be found in `demo.py`.
+
+### Demo #1:
+The following samples were generated with a fixed style and fixed bias.
+
+**Smash Mouth – All Star (<a href="https://www.azlyrics.com/lyrics/smashmouth/allstar.html">lyrics</a>)**
+![](img/all_star.svg)
+
+### Demo #2
+The following samples were generated with varying style and fixed bias.  Each verse is generated in a different style.
+
+**Vanessa Carlton – A Thousand Miles (<a href="https://www.azlyrics.com/lyrics/vanessacarlton/athousandmiles.html">lyrics</a>)**
+![](img/downtown.svg)
+
+### Demo #3
+The following samples were generated with a fixed style and varying bias.  Each verse has a lower bias than the previous, with the last verse being unbiased.
+
+**Leonard Cohen – Hallelujah (<a href="https://www.youtube.com/watch?v=dQw4w9WgXcQ">lyrics</a>)**
+![](img/give_up.svg)
+
+## Contribute
+This project was intended to serve as a reference implementation for a research paper, but since the results are of decent quality, it may be worthwile to make the project more broadly usable.  I plan to continue focusing on the machine learning side of things.  That said, I'd welcome contributors who can:
+
+  - Package this, and otherwise make it look more like a usable software project and less like research code.
+  - Add support for more sophisticated drawing, animations, or anything else in this direction.  Currently, the project only creates some simple svg files.

requirements.txt

@@ -0,0 +1,15 @@
+matplotlib==3.9.4
+pandas==2.2.3
+scikit-learn==1.6.1
+scipy==1.13.0
+svgwrite==1.1.12
+tensorflow-probability==0.23.0
+tensorflow==2.15.0
+fastapi==0.109.0
+uvicorn==0.27.0
+pydantic==2.5.0
+cairosvg==2.7.1
+gradio==4.29.0
+python-multipart==0.0.9
+httpx==0.27.0
+requests==2.31.0

rnn.py

@@ -0,0 +1,236 @@
+from __future__ import print_function
+import os
+
+import numpy as np
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+
+import drawing
+from data_frame import DataFrame
+from rnn_cell import LSTMAttentionCell
+from rnn_ops import rnn_free_run
+from tf_base_model import TFBaseModel
+from tf_utils import time_distributed_dense_layer
+
+
+class DataReader(object):
+
+    def __init__(self, data_dir):
+        data_cols = ['x', 'x_len', 'c', 'c_len']
+        data = [np.load(os.path.join(data_dir, '{}.npy'.format(i))) for i in data_cols]
+
+        self.test_df = DataFrame(columns=data_cols, data=data)
+        self.train_df, self.val_df = self.test_df.train_test_split(train_size=0.95, random_state=2018)
+
+        print('train size', len(self.train_df))
+        print('val size', len(self.val_df))
+        print('test size', len(self.test_df))
+
+    def train_batch_generator(self, batch_size):
+        return self.batch_generator(
+            batch_size=batch_size,
+            df=self.train_df,
+            shuffle=True,
+            num_epochs=10000,
+            mode='train'
+        )
+
+    def val_batch_generator(self, batch_size):
+        return self.batch_generator(
+            batch_size=batch_size,
+            df=self.val_df,
+            shuffle=True,
+            num_epochs=10000,
+            mode='val'
+        )
+
+    def test_batch_generator(self, batch_size):
+        return self.batch_generator(
+            batch_size=batch_size,
+            df=self.test_df,
+            shuffle=False,
+            num_epochs=1,
+            mode='test'
+        )
+
+    def batch_generator(self, batch_size, df, shuffle=True, num_epochs=10000, mode='train'):
+        gen = df.batch_generator(
+            batch_size=batch_size,
+            shuffle=shuffle,
+            num_epochs=num_epochs,
+            allow_smaller_final_batch=(mode == 'test')
+        )
+        for batch in gen:
+            batch['x_len'] = batch['x_len'] - 1
+            max_x_len = np.max(batch['x_len'])
+            max_c_len = np.max(batch['c_len'])
+            batch['y'] = batch['x'][:, 1:max_x_len + 1, :]
+            batch['x'] = batch['x'][:, :max_x_len, :]
+            batch['c'] = batch['c'][:, :max_c_len]
+            yield batch
+
+
+class rnn(TFBaseModel):
+
+    def __init__(
+        self,
+        lstm_size,
+        output_mixture_components,
+        attention_mixture_components,
+        **kwargs
+    ):
+        self.lstm_size = lstm_size
+        self.output_mixture_components = output_mixture_components
+        self.output_units = self.output_mixture_components*6 + 1
+        self.attention_mixture_components = attention_mixture_components
+        super(rnn, self).__init__(**kwargs)
+
+    def parse_parameters(self, z, eps=1e-8, sigma_eps=1e-4):
+        pis, sigmas, rhos, mus, es = tf.split(
+            z,
+            [
+                1*self.output_mixture_components,
+                2*self.output_mixture_components,
+                1*self.output_mixture_components,
+                2*self.output_mixture_components,
+                1
+            ],
+            axis=-1
+        )
+        pis = tf.nn.softmax(pis, axis=-1)
+        sigmas = tf.clip_by_value(tf.exp(sigmas), sigma_eps, np.inf)
+        rhos = tf.clip_by_value(tf.tanh(rhos), eps - 1.0, 1.0 - eps)
+        es = tf.clip_by_value(tf.nn.sigmoid(es), eps, 1.0 - eps)
+        return pis, mus, sigmas, rhos, es
+
+    def NLL(self, y, lengths, pis, mus, sigmas, rho, es, eps=1e-8):
+        sigma_1, sigma_2 = tf.split(sigmas, 2, axis=2)
+        y_1, y_2, y_3 = tf.split(y, 3, axis=2)
+        mu_1, mu_2 = tf.split(mus, 2, axis=2)
+
+        norm = 1.0 / (2*np.pi*sigma_1*sigma_2 * tf.sqrt(1 - tf.square(rho)))
+        Z = tf.square((y_1 - mu_1) / (sigma_1)) + \
+            tf.square((y_2 - mu_2) / (sigma_2)) - \
+            2*rho*(y_1 - mu_1)*(y_2 - mu_2) / (sigma_1*sigma_2)
+
+        exp = -1.0*Z / (2*(1 - tf.square(rho)))
+        gaussian_likelihoods = tf.exp(exp) * norm
+        gmm_likelihood = tf.reduce_sum(pis * gaussian_likelihoods, 2)
+        gmm_likelihood = tf.clip_by_value(gmm_likelihood, eps, np.inf)
+
+        bernoulli_likelihood = tf.squeeze(tf.where(tf.equal(tf.ones_like(y_3), y_3), es, 1 - es))
+
+        nll = -(tf.log(gmm_likelihood) + tf.log(bernoulli_likelihood))
+        sequence_mask = tf.logical_and(
+            tf.sequence_mask(lengths, maxlen=tf.shape(y)[1]),
+            tf.logical_not(tf.is_nan(nll)),
+        )
+        nll = tf.where(sequence_mask, nll, tf.zeros_like(nll))
+        num_valid = tf.reduce_sum(tf.cast(sequence_mask, tf.float32), axis=1)
+
+        sequence_loss = tf.reduce_sum(nll, axis=1) / tf.maximum(num_valid, 1.0)
+        element_loss = tf.reduce_sum(nll) / tf.maximum(tf.reduce_sum(num_valid), 1.0)
+        return sequence_loss, element_loss
+
+    def sample(self, cell):
+        initial_state = cell.zero_state(self.num_samples, dtype=tf.float32)
+        initial_input = tf.concat([
+            tf.zeros([self.num_samples, 2]),
+            tf.ones([self.num_samples, 1]),
+        ], axis=1)
+        return rnn_free_run(
+            cell=cell,
+            sequence_length=self.sample_tsteps,
+            initial_state=initial_state,
+            initial_input=initial_input,
+            scope='rnn'
+        )[1]
+
+    def primed_sample(self, cell):
+        initial_state = cell.zero_state(self.num_samples, dtype=tf.float32)
+        primed_state = tf.nn.dynamic_rnn(
+            inputs=self.x_prime,
+            cell=cell,
+            sequence_length=self.x_prime_len,
+            dtype=tf.float32,
+            initial_state=initial_state,
+            scope='rnn'
+        )[1]
+        return rnn_free_run(
+            cell=cell,
+            sequence_length=self.sample_tsteps,
+            initial_state=primed_state,
+            scope='rnn'
+        )[1]
+
+    def calculate_loss(self):
+        self.x = tf.placeholder(tf.float32, [None, None, 3])
+        self.y = tf.placeholder(tf.float32, [None, None, 3])
+        self.x_len = tf.placeholder(tf.int32, [None])
+        self.c = tf.placeholder(tf.int32, [None, None])
+        self.c_len = tf.placeholder(tf.int32, [None])
+
+        self.sample_tsteps = tf.placeholder(tf.int32, [])
+        self.num_samples = tf.placeholder(tf.int32, [])
+        self.prime = tf.placeholder(tf.bool, [])
+        self.x_prime = tf.placeholder(tf.float32, [None, None, 3])
+        self.x_prime_len = tf.placeholder(tf.int32, [None])
+        self.bias = tf.placeholder_with_default(
+            tf.zeros([self.num_samples], dtype=tf.float32), [None])
+
+        cell = LSTMAttentionCell(
+            lstm_size=self.lstm_size,
+            num_attn_mixture_components=self.attention_mixture_components,
+            attention_values=tf.one_hot(self.c, len(drawing.alphabet)),
+            attention_values_lengths=self.c_len,
+            num_output_mixture_components=self.output_mixture_components,
+            bias=self.bias
+        )
+        self.initial_state = cell.zero_state(tf.shape(self.x)[0], dtype=tf.float32)
+        outputs, self.final_state = tf.nn.dynamic_rnn(
+            inputs=self.x,
+            cell=cell,
+            sequence_length=self.x_len,
+            dtype=tf.float32,
+            initial_state=self.initial_state,
+            scope='rnn'
+        )
+        params = time_distributed_dense_layer(outputs, self.output_units, scope='rnn/gmm')
+        pis, mus, sigmas, rhos, es = self.parse_parameters(params)
+        sequence_loss, self.loss = self.NLL(self.y, self.x_len, pis, mus, sigmas, rhos, es)
+
+        self.sampled_sequence = tf.cond(
+            self.prime,
+            lambda: self.primed_sample(cell),
+            lambda: self.sample(cell)
+        )
+        return self.loss
+
+
+if __name__ == '__main__':
+    dr = DataReader(data_dir='data/processed/')
+
+    nn = rnn(
+        reader=dr,
+        log_dir='logs',
+        checkpoint_dir='checkpoints',
+        prediction_dir='predictions',
+        learning_rates=[.0001, .00005, .00002],
+        batch_sizes=[32, 64, 64],
+        patiences=[1500, 1000, 500],
+        beta1_decays=[.9, .9, .9],
+        validation_batch_size=32,
+        optimizer='rms',
+        num_training_steps=100000,
+        warm_start_init_step=0,
+        regularization_constant=0.0,
+        keep_prob=1.0,
+        enable_parameter_averaging=False,
+        min_steps_to_checkpoint=2000,
+        log_interval=20,
+        grad_clip=10,
+        lstm_size=400,
+        output_mixture_components=20,
+        attention_mixture_components=10
+    )
+    nn.fit()

rnn_cell.py

@@ -0,0 +1,185 @@
+from collections import namedtuple
+
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+import tensorflow_probability as tfp
+tfd = tfp.distributions
+import numpy as np
+
+from tf_utils import dense_layer, shape
+
+
+LSTMAttentionCellState = namedtuple(
+    'LSTMAttentionCellState',
+    ['h1', 'c1', 'h2', 'c2', 'h3', 'c3', 'alpha', 'beta', 'kappa', 'w', 'phi']
+)
+
+
+class LSTMAttentionCell(tf.nn.rnn_cell.RNNCell):
+
+    def __init__(
+        self,
+        lstm_size,
+        num_attn_mixture_components,
+        attention_values,
+        attention_values_lengths,
+        num_output_mixture_components,
+        bias,
+        reuse=None,
+    ):
+        self.reuse = reuse
+        self.lstm_size = lstm_size
+        self.num_attn_mixture_components = num_attn_mixture_components
+        self.attention_values = attention_values
+        self.attention_values_lengths = attention_values_lengths
+        self.window_size = shape(self.attention_values, 2)
+        self.char_len = tf.shape(attention_values)[1]
+        self.batch_size = tf.shape(attention_values)[0]
+        self.num_output_mixture_components = num_output_mixture_components
+        self.output_units = 6*self.num_output_mixture_components + 1
+        self.bias = bias
+
+    @property
+    def state_size(self):
+        return LSTMAttentionCellState(
+            self.lstm_size,
+            self.lstm_size,
+            self.lstm_size,
+            self.lstm_size,
+            self.lstm_size,
+            self.lstm_size,
+            self.num_attn_mixture_components,
+            self.num_attn_mixture_components,
+            self.num_attn_mixture_components,
+            self.window_size,
+            self.char_len,
+        )
+
+    @property
+    def output_size(self):
+        return self.lstm_size
+
+    def zero_state(self, batch_size, dtype):
+        return LSTMAttentionCellState(
+            tf.zeros([batch_size, self.lstm_size]),
+            tf.zeros([batch_size, self.lstm_size]),
+            tf.zeros([batch_size, self.lstm_size]),
+            tf.zeros([batch_size, self.lstm_size]),
+            tf.zeros([batch_size, self.lstm_size]),
+            tf.zeros([batch_size, self.lstm_size]),
+            tf.zeros([batch_size, self.num_attn_mixture_components]),
+            tf.zeros([batch_size, self.num_attn_mixture_components]),
+            tf.zeros([batch_size, self.num_attn_mixture_components]),
+            tf.zeros([batch_size, self.window_size]),
+            tf.zeros([batch_size, self.char_len]),
+        )
+
+    def __call__(self, inputs, state, scope=None):
+        with tf.variable_scope(scope or type(self).__name__, reuse=tf.AUTO_REUSE):
+
+            # lstm 1
+            s1_in = tf.concat([state.w, inputs], axis=1)
+            cell1 = tf.compat.v1.nn.rnn_cell.LSTMCell(self.lstm_size)
+            s1_out, s1_state = cell1(s1_in, state=(state.c1, state.h1))
+
+            # attention
+            attention_inputs = tf.concat([state.w, inputs, s1_out], axis=1)
+            attention_params = dense_layer(attention_inputs, 3*self.num_attn_mixture_components, scope='attention')
+            alpha, beta, kappa = tf.split(tf.nn.softplus(attention_params), 3, axis=1)
+            kappa = state.kappa + kappa / 25.0
+            beta = tf.clip_by_value(beta, .01, np.inf)
+
+            kappa_flat, alpha_flat, beta_flat = kappa, alpha, beta
+            kappa, alpha, beta = tf.expand_dims(kappa, 2), tf.expand_dims(alpha, 2), tf.expand_dims(beta, 2)
+
+            enum = tf.reshape(tf.range(self.char_len), (1, 1, self.char_len))
+            u = tf.cast(tf.tile(enum, (self.batch_size, self.num_attn_mixture_components, 1)), tf.float32)
+            phi_flat = tf.reduce_sum(alpha*tf.exp(-tf.square(kappa - u) / beta), axis=1)
+
+            phi = tf.expand_dims(phi_flat, 2)
+            sequence_mask = tf.cast(tf.sequence_mask(self.attention_values_lengths, maxlen=self.char_len), tf.float32)
+            sequence_mask = tf.expand_dims(sequence_mask, 2)
+            w = tf.reduce_sum(phi*self.attention_values*sequence_mask, axis=1)
+
+            # lstm 2
+            s2_in = tf.concat([inputs, s1_out, w], axis=1)
+            cell2 = tf.compat.v1.nn.rnn_cell.LSTMCell(self.lstm_size)
+            s2_out, s2_state = cell2(s2_in, state=(state.c2, state.h2))
+
+            # lstm 3
+            s3_in = tf.concat([inputs, s2_out, w], axis=1)
+            cell3 = tf.compat.v1.nn.rnn_cell.LSTMCell(self.lstm_size)
+            s3_out, s3_state = cell3(s3_in, state=(state.c3, state.h3))
+
+            new_state = LSTMAttentionCellState(
+                s1_state.h,
+                s1_state.c,
+                s2_state.h,
+                s2_state.c,
+                s3_state.h,
+                s3_state.c,
+                alpha_flat,
+                beta_flat,
+                kappa_flat,
+                w,
+                phi_flat,
+            )
+
+            return s3_out, new_state
+
+    def output_function(self, state):
+        params = dense_layer(state.h3, self.output_units, scope='gmm', reuse=tf.AUTO_REUSE)
+        pis, mus, sigmas, rhos, es = self._parse_parameters(params)
+        mu1, mu2 = tf.split(mus, 2, axis=1)
+        mus = tf.stack([mu1, mu2], axis=2)
+        sigma1, sigma2 = tf.split(sigmas, 2, axis=1)
+
+        covar_matrix = [tf.square(sigma1), rhos*sigma1*sigma2,
+                        rhos*sigma1*sigma2, tf.square(sigma2)]
+        covar_matrix = tf.stack(covar_matrix, axis=2)
+        covar_matrix = tf.reshape(covar_matrix, (self.batch_size, self.num_output_mixture_components, 2, 2))
+
+        mvn = tfd.MultivariateNormalFullCovariance(loc=mus, covariance_matrix=covar_matrix)
+        b = tfd.Bernoulli(probs=es)
+        c = tfd.Categorical(probs=pis)
+
+        sampled_e = b.sample()
+        sampled_coords = mvn.sample()
+        sampled_idx = c.sample()
+
+        idx = tf.stack([tf.range(self.batch_size), sampled_idx], axis=1)
+        coords = tf.gather_nd(sampled_coords, idx)
+        return tf.concat([coords, tf.cast(sampled_e, tf.float32)], axis=1)
+
+    def termination_condition(self, state):
+        char_idx = tf.cast(tf.argmax(state.phi, axis=1), tf.int32)
+        final_char = char_idx >= self.attention_values_lengths - 1
+        past_final_char = char_idx >= self.attention_values_lengths
+        output = self.output_function(state)
+        es = tf.cast(output[:, 2], tf.int32)
+        is_eos = tf.equal(es, tf.ones_like(es))
+        return tf.logical_or(tf.logical_and(final_char, is_eos), past_final_char)
+
+    def _parse_parameters(self, gmm_params, eps=1e-8, sigma_eps=1e-4):
+        pis, sigmas, rhos, mus, es = tf.split(
+            gmm_params,
+            [
+                1*self.num_output_mixture_components,
+                2*self.num_output_mixture_components,
+                1*self.num_output_mixture_components,
+                2*self.num_output_mixture_components,
+                1
+            ],
+            axis=-1
+        )
+        pis = pis*(1 + tf.expand_dims(self.bias, 1))
+        sigmas = sigmas - tf.expand_dims(self.bias, 1)
+
+        pis = tf.nn.softmax(pis, axis=-1)
+        pis = tf.where(pis < .01, tf.zeros_like(pis), pis)
+        sigmas = tf.clip_by_value(tf.exp(sigmas), sigma_eps, np.inf)
+        rhos = tf.clip_by_value(tf.tanh(rhos), eps - 1.0, 1.0 - eps)
+        es = tf.clip_by_value(tf.nn.sigmoid(es), eps, 1.0 - eps)
+        es = tf.where(es < .01, tf.zeros_like(es), es)
+
+        return pis, mus, sigmas, rhos, es

rnn_ops.py

@@ -0,0 +1,249 @@
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import cond as control_flow_ops_cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import variable_scope as vs
+from tensorflow.python.ops.rnn_cell_impl import _concat, assert_like_rnncell
+from tensorflow.python.ops.rnn import _maybe_tensor_shape_from_tensor
+from tensorflow.python.util import nest
+from tensorflow.python.framework import tensor_shape
+
+
+
+def raw_rnn(cell, loop_fn, parallel_iterations=None, swap_memory=False, scope=None):
+    """
+    raw_rnn adapted from the original tensorflow implementation
+    (https://github.com/tensorflow/tensorflow/blob/r1.4/tensorflow/python/ops/rnn.py)
+    to emit arbitrarily nested states for each time step (concatenated along the time axis)
+    in addition to the outputs at each timestep and the final state
+
+    returns (
+        states for all timesteps,
+        outputs for all timesteps,
+        final cell state,
+    )
+    """
+    assert_like_rnncell("dummy_name", cell)
+    if not callable(loop_fn):
+        raise TypeError("loop_fn must be a callable")
+
+    parallel_iterations = parallel_iterations or 32
+
+    # Create a new scope in which the caching device is either
+    # determined by the parent scope, or is set to place the cached
+    # Variable using the same placement as for the rest of the RNN.
+    with vs.variable_scope(scope or "rnn") as varscope:
+        if not tf.executing_eagerly():
+            if varscope.caching_device is None:
+                varscope.set_caching_device(lambda op: op.device)
+
+        time = constant_op.constant(0, dtype=dtypes.int32)
+        (elements_finished, next_input, initial_state, emit_structure,
+         init_loop_state) = loop_fn(time, None, None, None)
+        flat_input = nest.flatten(next_input)
+
+        # Need a surrogate loop state for the while_loop if none is available.
+        loop_state = (init_loop_state if init_loop_state is not None
+                      else constant_op.constant(0, dtype=dtypes.int32))
+
+        input_shape = [input_.get_shape() for input_ in flat_input]
+        static_batch_size = input_shape[0][0]
+
+        for input_shape_i in input_shape:
+            # Static verification that batch sizes all match
+            static_batch_size.merge_with(input_shape_i[0])
+
+        batch_size = static_batch_size.value
+        const_batch_size = batch_size
+        if batch_size is None:
+            batch_size = array_ops.shape(flat_input[0])[0]
+
+        nest.assert_same_structure(initial_state, cell.state_size)
+        state = initial_state
+        flat_state = nest.flatten(state)
+        flat_state = [ops.convert_to_tensor(s) for s in flat_state]
+        state = nest.pack_sequence_as(structure=state,
+                                      flat_sequence=flat_state)
+
+        if emit_structure is not None:
+            flat_emit_structure = nest.flatten(emit_structure)
+            flat_emit_size = [emit.shape if emit.shape.is_fully_defined() else
+                              array_ops.shape(emit) for emit in flat_emit_structure]
+            flat_emit_dtypes = [emit.dtype for emit in flat_emit_structure]
+        else:
+            emit_structure = cell.output_size
+            flat_emit_size = nest.flatten(emit_structure)
+            flat_emit_dtypes = [flat_state[0].dtype] * len(flat_emit_size)
+
+        flat_state_size = [s.shape if s.shape.is_fully_defined() else
+                           array_ops.shape(s) for s in flat_state]
+        flat_state_dtypes = [s.dtype for s in flat_state]
+
+        flat_emit_ta = [
+            tensor_array_ops.TensorArray(
+                dtype=dtype_i,
+                dynamic_size=True,
+                element_shape=(tensor_shape.TensorShape([const_batch_size])
+                               .concatenate(_maybe_tensor_shape_from_tensor(size_i))),
+                size=0,
+                name="rnn_output_%d" % i
+            )
+            for i, (dtype_i, size_i) in enumerate(zip(flat_emit_dtypes, flat_emit_size))
+        ]
+        emit_ta = nest.pack_sequence_as(structure=emit_structure, flat_sequence=flat_emit_ta)
+        flat_zero_emit = [
+            array_ops.zeros(_concat(batch_size, size_i), dtype_i)
+            for size_i, dtype_i in zip(flat_emit_size, flat_emit_dtypes)]
+
+        zero_emit = nest.pack_sequence_as(structure=emit_structure, flat_sequence=flat_zero_emit)
+
+        flat_state_ta = [
+            tensor_array_ops.TensorArray(
+                dtype=dtype_i,
+                dynamic_size=True,
+                element_shape=(tensor_shape.TensorShape([const_batch_size])
+                               .concatenate(_maybe_tensor_shape_from_tensor(size_i))),
+                size=0,
+                name="rnn_state_%d" % i
+            )
+            for i, (dtype_i, size_i) in enumerate(zip(flat_state_dtypes, flat_state_size))
+        ]
+        state_ta = nest.pack_sequence_as(structure=state, flat_sequence=flat_state_ta)
+
+        def condition(unused_time, elements_finished, *_):
+            return math_ops.logical_not(math_ops.reduce_all(elements_finished))
+
+        def body(time, elements_finished, current_input, state_ta, emit_ta, state, loop_state):
+            (next_output, cell_state) = cell(current_input, state)
+
+            nest.assert_same_structure(state, cell_state)
+            nest.assert_same_structure(cell.output_size, next_output)
+
+            next_time = time + 1
+            (next_finished, next_input, next_state, emit_output,
+             next_loop_state) = loop_fn(next_time, next_output, cell_state, loop_state)
+
+            nest.assert_same_structure(state, next_state)
+            nest.assert_same_structure(current_input, next_input)
+            nest.assert_same_structure(emit_ta, emit_output)
+
+            # If loop_fn returns None for next_loop_state, just reuse the previous one.
+            loop_state = loop_state if next_loop_state is None else next_loop_state
+
+            def _copy_some_through(current, candidate):
+                """Copy some tensors through via array_ops.where."""
+                def copy_fn(cur_i, cand_i):
+                    # TensorArray and scalar get passed through.
+                    if isinstance(cur_i, tensor_array_ops.TensorArray):
+                        return cand_i
+                    if cur_i.shape.ndims == 0:
+                        return cand_i
+                    # Otherwise propagate the old or the new value.
+                    with ops.colocate_with(cand_i):
+                        return array_ops.where(elements_finished, cur_i, cand_i)
+                return nest.map_structure(copy_fn, current, candidate)
+
+            emit_output = _copy_some_through(zero_emit, emit_output)
+            next_state = _copy_some_through(state, next_state)
+
+            emit_ta = nest.map_structure(lambda ta, emit: ta.write(time, emit), emit_ta, emit_output)
+            state_ta = nest.map_structure(lambda ta, state: ta.write(time, state), state_ta, next_state)
+
+            elements_finished = math_ops.logical_or(elements_finished, next_finished)
+
+            return (next_time, elements_finished, next_input, state_ta,
+                    emit_ta, next_state, loop_state)
+
+        returned = tf.while_loop(
+            condition, body, loop_vars=[
+                time, elements_finished, next_input, state_ta,
+                emit_ta, state, loop_state],
+            parallel_iterations=parallel_iterations,
+            swap_memory=swap_memory
+        )
+
+        (state_ta, emit_ta, final_state, final_loop_state) = returned[-4:]
+
+        flat_states = nest.flatten(state_ta)
+        flat_states = [array_ops.transpose(ta.stack(), (1, 0, 2)) for ta in flat_states]
+        states = nest.pack_sequence_as(structure=state_ta, flat_sequence=flat_states)
+
+        flat_outputs = nest.flatten(emit_ta)
+        flat_outputs = [array_ops.transpose(ta.stack(), (1, 0, 2)) for ta in flat_outputs]
+        outputs = nest.pack_sequence_as(structure=emit_ta, flat_sequence=flat_outputs)
+
+        return (states, outputs, final_state)
+
+
+def rnn_teacher_force(inputs, cell, sequence_length, initial_state, scope='dynamic-rnn-teacher-force'):
+    """
+    Implementation of an rnn with teacher forcing inputs provided.
+    Used in the same way as tf.dynamic_rnn.
+    """
+    inputs = array_ops.transpose(inputs, (1, 0, 2))
+    inputs_ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=array_ops.shape(inputs)[0])
+    inputs_ta = inputs_ta.unstack(inputs)
+
+    def loop_fn(time, cell_output, cell_state, loop_state):
+        emit_output = cell_output
+        next_cell_state = initial_state if cell_output is None else cell_state
+
+        elements_finished = time >= sequence_length
+        finished = math_ops.reduce_all(elements_finished)
+
+        next_input = control_flow_ops_cond.cond(
+            finished,
+            lambda: array_ops.zeros([array_ops.shape(inputs)[1], inputs.shape.as_list()[2]], dtype=dtypes.float32),
+            lambda: inputs_ta.read(time)
+        )
+
+        next_loop_state = None
+        return (elements_finished, next_input, next_cell_state, emit_output, next_loop_state)
+
+    states, outputs, final_state = raw_rnn(cell, loop_fn, scope=scope)
+    return states, outputs, final_state
+
+
+def rnn_free_run(cell, initial_state, sequence_length, initial_input=None, scope='dynamic-rnn-free-run'):
+    """
+    Implementation of an rnn which feeds its feeds its predictions back to itself at the next timestep.
+
+    cell must implement two methods:
+
+        cell.output_function(state) which takes in the state at timestep t and returns
+        the cell input at timestep t+1.
+
+        cell.termination_condition(state) which returns a boolean tensor of shape
+        [batch_size] denoting which sequences no longer need to be sampled.
+    """
+    with vs.variable_scope(scope, reuse=True):
+        if initial_input is None:
+            initial_input = cell.output_function(initial_state)
+
+    def loop_fn(time, cell_output, cell_state, loop_state):
+        next_cell_state = initial_state if cell_output is None else cell_state
+
+        elements_finished = math_ops.logical_or(
+            time >= sequence_length,
+            cell.termination_condition(next_cell_state)
+        )
+        finished = math_ops.reduce_all(elements_finished)
+
+        next_input = control_flow_ops_cond.cond(
+            finished,
+            lambda: array_ops.zeros_like(initial_input),
+            lambda: initial_input if cell_output is None else cell.output_function(next_cell_state)
+        )
+        emit_output = next_input[0] if cell_output is None else next_input
+
+        next_loop_state = None
+        return (elements_finished, next_input, next_cell_state, emit_output, next_loop_state)
+
+    states, outputs, final_state = raw_rnn(cell, loop_fn, scope=scope)
+    return states, outputs, final_state

tf_base_model.py

@@ -0,0 +1,408 @@
+from __future__ import print_function
+from collections import deque
+from datetime import datetime
+import logging
+import os
+import pprint as pp
+import time
+
+import numpy as np
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+
+from tf_utils import shape
+
+
+class TFBaseModel(object):
+
+    """Interface containing some boilerplate code for training tensorflow models.
+
+    Subclassing models must implement self.calculate_loss(), which returns a tensor for the batch loss.
+    Code for the training loop, parameter updates, checkpointing, and inference are implemented here and
+    subclasses are mainly responsible for building the computational graph beginning with the placeholders
+    and ending with the loss tensor.
+
+    Args:
+        reader: Class with attributes train_batch_generator, val_batch_generator, and test_batch_generator
+            that yield dictionaries mapping tf.placeholder names (as strings) to batch data (numpy arrays).
+        batch_size: Minibatch size.
+        learning_rate: Learning rate.
+        optimizer: 'rms' for RMSProp, 'adam' for Adam, 'sgd' for SGD
+        grad_clip: Clip gradients elementwise to have norm at most equal to grad_clip.
+        regularization_constant:  Regularization constant applied to all trainable parameters.
+        keep_prob: 1 - p, where p is the dropout probability
+        early_stopping_steps:  Number of steps to continue training after validation loss has
+            stopped decreasing.
+        warm_start_init_step:  If nonzero, model will resume training a restored model beginning
+            at warm_start_init_step.
+        num_restarts:  After validation loss plateaus, the best checkpoint will be restored and the
+            learning rate will be halved.  This process will repeat num_restarts times.
+        enable_parameter_averaging:  If true, model saves exponential weighted averages of parameters
+            to separate checkpoint file.
+        min_steps_to_checkpoint:  Model only saves after min_steps_to_checkpoint training steps
+            have passed.
+        log_interval:  Train and validation accuracies are logged every log_interval training steps.
+        loss_averaging_window:  Train/validation losses are averaged over the last loss_averaging_window
+            training steps.
+        num_validation_batches:  Number of batches to be used in validation evaluation at each step.
+        log_dir: Directory where logs are written.
+        checkpoint_dir: Directory where checkpoints are saved.
+        prediction_dir: Directory where predictions/outputs are saved.
+    """
+
+    def __init__(
+        self,
+        reader=None,
+        batch_sizes=[128],
+        num_training_steps=20000,
+        learning_rates=[.01],
+        beta1_decays=[.99],
+        optimizer='adam',
+        grad_clip=5,
+        regularization_constant=0.0,
+        keep_prob=1.0,
+        patiences=[3000],
+        warm_start_init_step=0,
+        enable_parameter_averaging=False,
+        min_steps_to_checkpoint=100,
+        log_interval=20,
+        logging_level=logging.INFO,
+        loss_averaging_window=100,
+        validation_batch_size=64,
+        log_dir='logs',
+        checkpoint_dir='checkpoints',
+        prediction_dir='predictions',
+    ):
+
+        assert len(batch_sizes) == len(learning_rates) == len(patiences)
+        self.batch_sizes = batch_sizes
+        self.learning_rates = learning_rates
+        self.beta1_decays = beta1_decays
+        self.patiences = patiences
+        self.num_restarts = len(batch_sizes) - 1
+        self.restart_idx = 0
+        self.update_train_params()
+
+        self.reader = reader
+        self.num_training_steps = num_training_steps
+        self.optimizer = optimizer
+        self.grad_clip = grad_clip
+        self.regularization_constant = regularization_constant
+        self.warm_start_init_step = warm_start_init_step
+        self.keep_prob_scalar = keep_prob
+        self.enable_parameter_averaging = enable_parameter_averaging
+        self.min_steps_to_checkpoint = min_steps_to_checkpoint
+        self.log_interval = log_interval
+        self.loss_averaging_window = loss_averaging_window
+        self.validation_batch_size = validation_batch_size
+
+        self.log_dir = log_dir
+        self.logging_level = logging_level
+        self.prediction_dir = prediction_dir
+        self.checkpoint_dir = checkpoint_dir
+        if self.enable_parameter_averaging:
+            self.checkpoint_dir_averaged = checkpoint_dir + '_avg'
+
+        self.init_logging(self.log_dir)
+        logging.info('\nnew run with parameters:\n{}'.format(pp.pformat(self.__dict__)))
+
+        self.graph = self.build_graph()
+        self.session = tf.Session(graph=self.graph)
+        logging.info('built graph')
+
+    def update_train_params(self):
+        self.batch_size = self.batch_sizes[self.restart_idx]
+        self.learning_rate = self.learning_rates[self.restart_idx]
+        self.beta1_decay = self.beta1_decays[self.restart_idx]
+        self.early_stopping_steps = self.patiences[self.restart_idx]
+
+    def calculate_loss(self):
+        raise NotImplementedError('subclass must implement this')
+
+    def fit(self):
+        with self.session.as_default():
+
+            if self.warm_start_init_step:
+                self.restore(self.warm_start_init_step)
+                step = self.warm_start_init_step
+            else:
+                self.session.run(self.init)
+                step = 0
+
+            train_generator = self.reader.train_batch_generator(self.batch_size)
+            val_generator = self.reader.val_batch_generator(self.validation_batch_size)
+
+            train_loss_history = deque(maxlen=self.loss_averaging_window)
+            val_loss_history = deque(maxlen=self.loss_averaging_window)
+            train_time_history = deque(maxlen=self.loss_averaging_window)
+            val_time_history = deque(maxlen=self.loss_averaging_window)
+            if not hasattr(self, 'metrics'):
+                self.metrics = {}
+
+            metric_histories = {
+                metric_name: deque(maxlen=self.loss_averaging_window) for metric_name in self.metrics
+            }
+            best_validation_loss, best_validation_tstep = float('inf'), 0
+
+            while step < self.num_training_steps:
+
+                # validation evaluation
+                val_start = time.time()
+                val_batch_df = next(val_generator)
+                val_feed_dict = {
+                    getattr(self, placeholder_name, None): data
+                    for placeholder_name, data in val_batch_df.items() if hasattr(self, placeholder_name)
+                }
+
+                val_feed_dict.update({self.learning_rate_var: self.learning_rate, self.beta1_decay_var: self.beta1_decay})
+                if hasattr(self, 'keep_prob'):
+                    val_feed_dict.update({self.keep_prob: 1.0})
+                if hasattr(self, 'is_training'):
+                    val_feed_dict.update({self.is_training: False})
+
+                results = self.session.run(
+                    fetches=[self.loss] + self.metrics.values(),
+                    feed_dict=val_feed_dict
+                )
+                val_loss = results[0]
+                val_metrics = results[1:] if len(results) > 1 else []
+                val_metrics = dict(zip(self.metrics.keys(), val_metrics))
+                val_loss_history.append(val_loss)
+                val_time_history.append(time.time() - val_start)
+                for key in val_metrics:
+                    metric_histories[key].append(val_metrics[key])
+
+                if hasattr(self, 'monitor_tensors'):
+                    for name, tensor in self.monitor_tensors.items():
+                        [np_val] = self.session.run([tensor], feed_dict=val_feed_dict)
+                        print(name)
+                        print('min', np_val.min())
+                        print('max', np_val.max())
+                        print('mean', np_val.mean())
+                        print('std', np_val.std())
+                        print('nans', np.isnan(np_val).sum())
+                        print()
+                    print()
+                    print()
+
+                # train step
+                train_start = time.time()
+                train_batch_df = next(train_generator)
+                train_feed_dict = {
+                    getattr(self, placeholder_name, None): data
+                    for placeholder_name, data in train_batch_df.items() if hasattr(self, placeholder_name)
+                }
+
+                train_feed_dict.update({self.learning_rate_var: self.learning_rate, self.beta1_decay_var: self.beta1_decay})
+                if hasattr(self, 'keep_prob'):
+                    train_feed_dict.update({self.keep_prob: self.keep_prob_scalar})
+                if hasattr(self, 'is_training'):
+                    train_feed_dict.update({self.is_training: True})
+
+                train_loss, _ = self.session.run(
+                    fetches=[self.loss, self.step],
+                    feed_dict=train_feed_dict
+                )
+                train_loss_history.append(train_loss)
+                train_time_history.append(time.time() - train_start)
+
+                if step % self.log_interval == 0:
+                    avg_train_loss = sum(train_loss_history) / len(train_loss_history)
+                    avg_val_loss = sum(val_loss_history) / len(val_loss_history)
+                    avg_train_time = sum(train_time_history) / len(train_time_history)
+                    avg_val_time = sum(val_time_history) / len(val_time_history)
+                    metric_log = (
+                        "[[step {:>8}]]     "
+                        "[[train {:>4}s]]     loss: {:<12}     "
+                        "[[val {:>4}s]]     loss: {:<12}     "
+                    ).format(
+                        step,
+                        round(avg_train_time, 4),
+                        round(avg_train_loss, 8),
+                        round(avg_val_time, 4),
+                        round(avg_val_loss, 8),
+                    )
+                    early_stopping_metric = avg_val_loss
+                    for metric_name, metric_history in metric_histories.items():
+                        metric_val = sum(metric_history) / len(metric_history)
+                        metric_log += '{}: {:<4}     '.format(metric_name, round(metric_val, 4))
+                        if metric_name == self.early_stopping_metric:
+                            early_stopping_metric = metric_val
+
+                    logging.info(metric_log)
+
+                    if early_stopping_metric < best_validation_loss:
+                        best_validation_loss = early_stopping_metric
+                        best_validation_tstep = step
+                        if step > self.min_steps_to_checkpoint:
+                            self.save(step)
+                            if self.enable_parameter_averaging:
+                                self.save(step, averaged=True)
+
+                    if step - best_validation_tstep > self.early_stopping_steps:
+
+                        if self.num_restarts is None or self.restart_idx >= self.num_restarts:
+                            logging.info('best validation loss of {} at training step {}'.format(
+                                best_validation_loss, best_validation_tstep))
+                            logging.info('early stopping - ending training.')
+                            return
+
+                        if self.restart_idx < self.num_restarts:
+                            self.restore(best_validation_tstep)
+                            step = best_validation_tstep
+                            self.restart_idx += 1
+                            self.update_train_params()
+                            train_generator = self.reader.train_batch_generator(self.batch_size)
+
+                step += 1
+
+            if step <= self.min_steps_to_checkpoint:
+                best_validation_tstep = step
+                self.save(step)
+                if self.enable_parameter_averaging:
+                    self.save(step, averaged=True)
+
+            logging.info('num_training_steps reached - ending training')
+
+    def predict(self, chunk_size=256):
+        if not os.path.isdir(self.prediction_dir):
+            os.makedirs(self.prediction_dir)
+
+        if hasattr(self, 'prediction_tensors'):
+            prediction_dict = {tensor_name: [] for tensor_name in self.prediction_tensors}
+
+            test_generator = self.reader.test_batch_generator(chunk_size)
+            for i, test_batch_df in enumerate(test_generator):
+                if i % 10 == 0:
+                    print(i*len(test_batch_df))
+
+                test_feed_dict = {
+                    getattr(self, placeholder_name, None): data
+                    for placeholder_name, data in test_batch_df.items() if hasattr(self, placeholder_name)
+                }
+                if hasattr(self, 'keep_prob'):
+                    test_feed_dict.update({self.keep_prob: 1.0})
+                if hasattr(self, 'is_training'):
+                    test_feed_dict.update({self.is_training: False})
+
+                tensor_names, tf_tensors = zip(*self.prediction_tensors.items())
+                np_tensors = self.session.run(
+                    fetches=tf_tensors,
+                    feed_dict=test_feed_dict
+                )
+                for tensor_name, tensor in zip(tensor_names, np_tensors):
+                    prediction_dict[tensor_name].append(tensor)
+
+            for tensor_name, tensor in prediction_dict.items():
+                np_tensor = np.concatenate(tensor, 0)
+                save_file = os.path.join(self.prediction_dir, '{}.npy'.format(tensor_name))
+                logging.info('saving {} with shape {} to {}'.format(tensor_name, np_tensor.shape, save_file))
+                np.save(save_file, np_tensor)
+
+        if hasattr(self, 'parameter_tensors'):
+            for tensor_name, tensor in self.parameter_tensors.items():
+                np_tensor = tensor.eval(self.session)
+
+                save_file = os.path.join(self.prediction_dir, '{}.npy'.format(tensor_name))
+                logging.info('saving {} with shape {} to {}'.format(tensor_name, np_tensor.shape, save_file))
+                np.save(save_file, np_tensor)
+
+    def save(self, step, averaged=False):
+        saver = self.saver_averaged if averaged else self.saver
+        checkpoint_dir = self.checkpoint_dir_averaged if averaged else self.checkpoint_dir
+        if not os.path.isdir(checkpoint_dir):
+            logging.info('creating checkpoint directory {}'.format(checkpoint_dir))
+            os.mkdir(checkpoint_dir)
+
+        model_path = os.path.join(checkpoint_dir, 'model')
+        logging.info('saving model to {}'.format(model_path))
+        saver.save(self.session, model_path, global_step=step)
+
+    def restore(self, step=None, averaged=False):
+        saver = self.saver_averaged if averaged else self.saver
+        checkpoint_dir = self.checkpoint_dir_averaged if averaged else self.checkpoint_dir
+        if not step:
+            model_path = tf.train.latest_checkpoint(checkpoint_dir)
+            logging.info('restoring model parameters from {}'.format(model_path))
+            saver.restore(self.session, model_path)
+        else:
+            model_path = os.path.join(
+                checkpoint_dir, 'model{}-{}'.format('_avg' if averaged else '', step)
+            )
+            logging.info('restoring model from {}'.format(model_path))
+            saver.restore(self.session, model_path)
+
+    def init_logging(self, log_dir):
+        if not os.path.isdir(log_dir):
+            os.makedirs(log_dir)
+
+        date_str = datetime.now().strftime('%Y-%m-%d_%H-%M')
+        log_file = 'log_{}.txt'.format(date_str)
+
+        try:                 # Python 2
+            reload(logging)  # bad
+        except NameError:    # Python 3
+            import logging
+        logging.basicConfig(
+            filename=os.path.join(log_dir, log_file),
+            level=self.logging_level,
+            format='[[%(asctime)s]] %(message)s',
+            datefmt='%m/%d/%Y %I:%M:%S %p'
+        )
+        logging.getLogger().addHandler(logging.StreamHandler())
+
+    def update_parameters(self, loss):
+        if self.regularization_constant != 0:
+            l2_norm = tf.reduce_sum([tf.sqrt(tf.reduce_sum(tf.square(param))) for param in tf.trainable_variables()])
+            loss = loss + self.regularization_constant*l2_norm
+
+        optimizer = self.get_optimizer(self.learning_rate_var, self.beta1_decay_var)
+        grads = optimizer.compute_gradients(loss)
+        clipped = [(tf.clip_by_value(g, -self.grad_clip, self.grad_clip), v_) for g, v_ in grads]
+
+        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+        with tf.control_dependencies(update_ops):
+            step = optimizer.apply_gradients(clipped, global_step=self.global_step)
+
+        if self.enable_parameter_averaging:
+            maintain_averages_op = self.ema.apply(tf.trainable_variables())
+            with tf.control_dependencies([step]):
+                self.step = tf.group(maintain_averages_op)
+        else:
+            self.step = step
+
+        logging.info('all parameters:')
+        logging.info(pp.pformat([(var.name, shape(var)) for var in tf.global_variables()]))
+
+        logging.info('trainable parameters:')
+        logging.info(pp.pformat([(var.name, shape(var)) for var in tf.trainable_variables()]))
+
+        logging.info('trainable parameter count:')
+        logging.info(str(np.sum(np.prod(shape(var)) for var in tf.trainable_variables())))
+
+    def get_optimizer(self, learning_rate, beta1_decay):
+        if self.optimizer == 'adam':
+            return tf.train.AdamOptimizer(learning_rate, beta1=beta1_decay)
+        elif self.optimizer == 'gd':
+            return tf.train.GradientDescentOptimizer(learning_rate)
+        elif self.optimizer == 'rms':
+            return tf.train.RMSPropOptimizer(learning_rate, decay=beta1_decay, momentum=0.9)
+        else:
+            assert False, 'optimizer must be adam, gd, or rms'
+
+    def build_graph(self):
+        with tf.Graph().as_default() as graph:
+            self.ema = tf.train.ExponentialMovingAverage(decay=0.99)
+            self.global_step = tf.Variable(0, trainable=False)
+            self.learning_rate_var = tf.Variable(0.0, trainable=False)
+            self.beta1_decay_var = tf.Variable(0.0, trainable=False)
+
+            self.loss = self.calculate_loss()
+            self.update_parameters(self.loss)
+
+            self.saver = tf.train.Saver(max_to_keep=1)
+            if self.enable_parameter_averaging:
+                self.saver_averaged = tf.train.Saver(self.ema.variables_to_restore(), max_to_keep=1)
+
+            self.init = tf.global_variables_initializer()
+            return graph

tf_utils.py

@@ -0,0 +1,91 @@
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+
+
+def dense_layer(inputs, output_units, bias=True, activation=None, batch_norm=None,
+                dropout=None, scope='dense-layer', reuse=False):
+    """
+    Applies a dense layer to a 2D tensor of shape [batch_size, input_units]
+    to produce a tensor of shape [batch_size, output_units].
+    Args:
+        inputs: Tensor of shape [batch size, input_units].
+        output_units: Number of output units.
+        activation: activation function.
+        dropout: dropout keep prob.
+    Returns:
+        Tensor of shape [batch size, output_units].
+    """
+    with tf.variable_scope(scope, reuse=reuse):
+        W = tf.get_variable(
+            name='weights',
+            initializer=tf.compat.v1.variance_scaling_initializer(),
+            shape=[shape(inputs, -1), output_units]
+        )
+        z = tf.matmul(inputs, W)
+        if bias:
+            b = tf.get_variable(
+                name='biases',
+                initializer=tf.constant_initializer(),
+                shape=[output_units]
+            )
+            z = z + b
+
+        if batch_norm is not None:
+            z = tf.layers.batch_normalization(z, training=batch_norm, reuse=reuse)
+
+        z = activation(z) if activation else z
+        z = tf.nn.dropout(z, dropout) if dropout is not None else z
+        return z
+
+
+def time_distributed_dense_layer(
+        inputs, output_units, bias=True, activation=None, batch_norm=None,
+        dropout=None, scope='time-distributed-dense-layer', reuse=False):
+    """
+    Applies a shared dense layer to each timestep of a tensor of shape
+    [batch_size, max_seq_len, input_units] to produce a tensor of shape
+    [batch_size, max_seq_len, output_units].
+
+    Args:
+        inputs: Tensor of shape [batch size, max sequence length, ...].
+        output_units: Number of output units.
+        activation: activation function.
+        dropout: dropout keep prob.
+
+    Returns:
+        Tensor of shape [batch size, max sequence length, output_units].
+    """
+    with tf.variable_scope(scope, reuse=reuse):
+        W = tf.get_variable(
+            name='weights',
+            initializer=tf.compat.v1.variance_scaling_initializer(),
+            shape=[shape(inputs, -1), output_units]
+        )
+        z = tf.einsum('ijk,kl->ijl', inputs, W)
+        if bias:
+            b = tf.get_variable(
+                name='biases',
+                initializer=tf.constant_initializer(),
+                shape=[output_units]
+            )
+            z = z + b
+
+        if batch_norm is not None:
+            z = tf.layers.batch_normalization(z, training=batch_norm, reuse=reuse)
+
+        z = activation(z) if activation else z
+        z = tf.nn.dropout(z, dropout) if dropout is not None else z
+        return z
+
+
+def shape(tensor, dim=None):
+    """Get tensor shape/dimension as list/int"""
+    if dim is None:
+        return tensor.shape.as_list()
+    else:
+        return tensor.shape.as_list()[dim]
+
+
+def rank(tensor):
+    """Get tensor rank as python list"""
+    return len(tensor.shape.as_list())