├── .colorize
├── .gitignore
├── .gitmodules
├── README.md
├── TestEnergyFunction.ipynb
├── TestPendulumGeneration.ipynb
├── VisualizePendulumCache.ipynb
├── datasets
    ├── imagepairs.py
    ├── images.py
    └── pendulum.py
├── models
    ├── pendulum_energy.py
    ├── simple.py
    ├── stabledynamics.py
    ├── sympy2torch.py
    ├── train_many
    ├── vae.py
    └── vae_trajectory.py
├── pendulum_error.py
├── plot_data.py
├── render.py
├── render_vae
├── requirements.txt
├── test_lstm.py
├── train.py
├── train_lstm.py
├── train_many
├── train_pendulum_simple
├── train_pendulum_stable
├── train_pendulum_stable_knownV
├── train_vae_simple
├── train_vae_stable
├── trajectory.py
├── util.py
└── vae_trajmodel.py


/.colorize:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 | 
3 | ("$@" 2>&1>&3 | sed -e $'s,\[[0-9:]* INFO.*,\e[93m&\e[m,' -e $'s,\[[0-9:]* INFO @root].*,\e[36m&\e[m,' -e $'s,\[[0-9:]* WARNING.*,\e[91m&\e[m,' -e $'s,\[[0-9:]* ERROR.*,\e[41m&\e[m,' >&2) 3>&1 | sed $'s,.*,\e[m&\e[m,'
4 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
 1 | .python-version
 2 | .vscode/
 3 | *.log
 4 | 
 5 | youtube
 6 | experiments/
 7 | pendulum-cache/
 8 | runs/
 9 | 
10 | **/__pycache__/
11 | .ipynb_checkpoints
12 | 


--------------------------------------------------------------------------------
/.gitmodules:
--------------------------------------------------------------------------------
1 | [submodule "models/ritheshkumar95_vqvae"]
2 | 	path = models/ritheshkumar95_vqvae
3 | 	url = git@github.com:ritheshkumar95/pytorch-vqvae.git
4 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Learning Stable Deep Dynamics Models
 2 | 
 3 | Companion code to "Learning Stable Deep Dynamics Models" (Manek and Kolter, 2019)
 4 | 
 5 | ## Installation
 6 | 
 7 | You need Python 3.6 or later, with packages listed in `requirements.txt`. Ensure that the directories `experiments/` and `runs/` are writable.
 8 | 
 9 | ## Running
10 | 
11 | All training commands produce output in the `experiments/` folder, automatically named with the hyperparameters used in training. You can track training progress using `tensorboard runs/`.
12 | 
13 | ### Pendulum
14 | 
15 | To train the models for an `<n>`-link pendulum, you should first run the command:
16 | 
17 | ```sh
18 | ./train_pendulum_simple <n>
19 | ```
20 | 
21 | which will create and cache the training data and the evaluation data. (Warning: this is a cpu-heavy task that may take hours to complete.) Once that is complete, you can train multiple models in parallel with:
22 | 
23 | ```sh
24 | ./train_many <n>
25 | ```
26 | 
27 | You can concurrently train train models with different `<n>`.
28 | 
29 | Ensure that `pendulum-cache/` is writable.
30 | 
31 | ### VAE
32 | 
33 | To train the VAE, you need to convert videos to sequentially-numbered frames in `youtube/<video>/fr_[0-9]*.jpg`. A good place to start is [this](https://www.youtube.com/watch?v=I4dYM1biVq0) video, which you can download and convert with:
34 | 
35 | ```sh
36 | mkdir youtube
37 | cd youtube
38 | youtube-dl -o bonfire.%(ext) -f 134 I4dYM1biVq0
39 | ffmpeg -ss 0:01:00 -i bonfire.mp4 -t 3:20 -vf scale=320:240 bonfire/fr_%06d.png
40 | ```
41 | 
42 | (You need to install `youtube-dl` and `ffmpeg`.)
43 | 
44 | Once that is done, you can train the simple VAE model with:
45 | 
46 | ```sh
47 | ./train_vae_simple bonfire
48 | ```
49 | 
50 | And the stable VAE with:
51 | 
52 | ```sh
53 | ./train_vae_stable stable_exp bonfire 0.0005 -0.25 0.00 PSICNN 0.001
54 | ```
55 | 


--------------------------------------------------------------------------------
/datasets/imagepairs.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | from . import images
 4 | import torch
 5 | 
 6 | class SeqPairs(torch.utils.data.Dataset):
 7 |     def __init__(self, image_dataset):
 8 |         super().__init__()
 9 |         self.image_dataset = image_dataset
10 | 
11 |     def __getitem__(self, index):
12 |         q1, _ = self.image_dataset[index]
13 |         q2, _ = self.image_dataset[index + 1]
14 |         return ((q1, q2), (q1, q2))
15 | 
16 |     def __len__(self):
17 |         return len(self.image_dataset) - 1
18 | 
19 | def build(props):
20 |     return SeqPairs(images.build(props))
21 | 


--------------------------------------------------------------------------------
/datasets/images.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import glob
 4 | import logging
 5 | from pathlib import Path
 6 | 
 7 | from PIL import Image
 8 | 
 9 | import torch
10 | from torch.utils.data import DataLoader
11 | from torchvision import transforms
12 | 
13 | 
14 | def make_flat_dataset(dir, extensions):
15 |     images = []
16 | 
17 |     for root, _, fnames in sorted(os.walk(os.path.expanduser(dir))):
18 |         for fname in sorted(fnames):
19 |             _, ext = os.path.splitext(fname)
20 |             if ext in extensions:
21 |                 path = os.path.join(root, fname)
22 |                 images.append((path, 0))
23 | 
24 |     return images
25 | 
26 | 
27 | def pil_loader(path):
28 |     # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
29 |     with open(path, 'rb') as f:
30 |         img = Image.open(f)
31 |         return img.convert('RGB')
32 | 
33 | class Images(torch.utils.data.Dataset):
34 |     """
35 |     Args:
36 |         glob_files (string): glob pointing to the files.
37 |         transform (callable, optional): A function/transform that  takes in an PIL image
38 |             and returns a transformed version. E.g, ``transforms.RandomCrop``
39 |         target_transform (callable, optional): A function/transform that takes in the
40 |             target and transforms it.
41 |         loader (callable, optional): A function to load an image given its path.
42 |      Attributes:
43 |         classes (list): List of the class names.
44 |         class_to_idx (dict): Dict with items (class_name, class_index).
45 |         imgs (list): List of (image path, class_index) tuples
46 |     """
47 |     def __init__(self, glob_files, transform=transforms.Compose([transforms.ToTensor()]), cache=True):
48 |         super().__init__()
49 | 
50 |         samples = glob.glob(glob_files)
51 |         assert len(samples) > 0, "Found 0 files matching {}".format(glob_files)
52 | 
53 |         self.loader = pil_loader
54 |         self.samples = sorted(samples)
55 |         self.transform = transform
56 |         # In-memory cache of transformed data:
57 |         # Set to None to prevent caching. 
58 |         self.cache = [None for _ in samples] if cache else None
59 | 
60 |     def _get_from_cache(self, index):
61 |         if self.cache is None or self.cache[index] is None:
62 |             path_a = self.samples[index]
63 |             sample_a = self.loader(path_a)
64 |             if self.transform is not None:
65 |                 sample_a = self.transform(sample_a)
66 |             if self.cache is not None:
67 |                 self.cache[index] = sample_a
68 |         else:
69 |             sample_a = self.cache[index]
70 | 
71 |         return sample_a
72 | 
73 |     def __getitem__(self, index):
74 |         """
75 |         Args:
76 |             index (int): Index
77 |         Returns:
78 |             tuple: (sample, target) where target is class_index of the target class.
79 |         """
80 |         rv = self._get_from_cache(index)
81 |         return (rv, rv)
82 | 
83 |     def __len__(self):
84 |         return len(self.samples)
85 | 
86 | def build(props):
87 |     return Images(props["files"], cache="nocache" not in props)
88 | 


--------------------------------------------------------------------------------
/datasets/pendulum.py:
--------------------------------------------------------------------------------
  1 | from pathlib import Path
  2 | 
  3 | import matplotlib.pyplot as plt
  4 | import numpy as np
  5 | import torch
  6 | from scipy.integrate import odeint
  7 | from sympy import Dummy, lambdify, symbols
  8 | from sympy.physics import mechanics
  9 | 
 10 | CACHE = Path("pendulum-cache/")
 11 | 
 12 | # Modified from: http://jakevdp.github.io/blog/2017/03/08/triple-pendulum-chaos/
 13 | def pendulum_gradient(n, lengths=None, masses=1, friction=0.3):
 14 |     """Integrate a multi-pendulum with `n` sections"""
 15 |     #-------------------------------------------------
 16 |     # Step 1: construct the pendulum model
 17 |     
 18 |     # Generalized coordinates and velocities
 19 |     # (in this case, angular positions & velocities of each mass) 
 20 |     q = mechanics.dynamicsymbols('q:{0}'.format(n))
 21 |     u = mechanics.dynamicsymbols('u:{0}'.format(n))
 22 | 
 23 |     # mass and length
 24 |     m = symbols('m:{0}'.format(n))
 25 |     l = symbols('l:{0}'.format(n))
 26 | 
 27 |     # gravity and time symbols
 28 |     g, t = symbols('g,t')
 29 |     
 30 |     #--------------------------------------------------
 31 |     # Step 2: build the model using Kane's Method
 32 | 
 33 |     # Create pivot point reference frame
 34 |     A = mechanics.ReferenceFrame('A')
 35 |     P = mechanics.Point('P')
 36 |     P.set_vel(A, 0)
 37 | 
 38 |     # lists to hold particles, forces, and kinetic ODEs
 39 |     # for each pendulum in the chain
 40 |     particles = []
 41 |     forces = []
 42 |     kinetic_odes = []
 43 | 
 44 |     for i in range(n):
 45 |         # Create a reference frame following the i^th mass
 46 |         Ai = A.orientnew('A' + str(i), 'Axis', [q[i], A.z])
 47 |         Ai.set_ang_vel(A, u[i] * A.z)
 48 | 
 49 |         # Create a point in this reference frame
 50 |         Pi = P.locatenew('P' + str(i), l[i] * Ai.x)
 51 |         Pi.v2pt_theory(P, A, Ai)
 52 | 
 53 |         # Create a new particle of mass m[i] at this point
 54 |         Pai = mechanics.Particle('Pa' + str(i), Pi, m[i])
 55 |         particles.append(Pai)
 56 | 
 57 |         # Set forces & compute kinematic ODE
 58 |         forces.append((Pi, m[i] * g * A.x))
 59 |         # Add damping torque:
 60 |         forces.append((Ai, -1 * friction * u[i] * A.z))
 61 | 
 62 |         kinetic_odes.append(q[i].diff(t) - u[i])
 63 | 
 64 |         P = Pi
 65 | 
 66 |     # Generate equations of motion
 67 |     KM = mechanics.KanesMethod(A, q_ind=q, u_ind=u,
 68 |                                kd_eqs=kinetic_odes)
 69 |     fr, fr_star = KM.kanes_equations(particles, forces)
 70 | 
 71 |     #-----------------------------------------------------
 72 |     # Step 3: numerically evaluate equations and integrate
 73 | 
 74 |     # lengths and masses
 75 |     if lengths is None:
 76 |         lengths = np.ones(n) / n
 77 |     lengths = np.broadcast_to(lengths, n)
 78 |     masses = np.broadcast_to(masses, n)
 79 | 
 80 |     # Fixed parameters: gravitational constant, lengths, and masses
 81 |     parameters = [g] + list(l) + list(m)
 82 |     parameter_vals = [9.81] + list(lengths) + list(masses)
 83 | 
 84 |     # define symbols for unknown parameters
 85 |     unknowns = [Dummy() for i in q + u]
 86 |     unknown_dict = dict(zip(q + u, unknowns))
 87 |     kds = KM.kindiffdict()
 88 | 
 89 |     # substitute unknown symbols for qdot terms
 90 |     mm_sym = KM.mass_matrix_full.subs(kds).subs(unknown_dict)
 91 |     fo_sym = KM.forcing_full.subs(kds).subs(unknown_dict)
 92 | 
 93 |     # create functions for numerical calculation
 94 |     mm_func = lambdify(unknowns + parameters, mm_sym)
 95 |     fo_func = lambdify(unknowns + parameters, fo_sym)
 96 | 
 97 |     # function which computes the derivatives of parameters
 98 |     def gradient(y, *a, **kw):
 99 |         squeeze = False
100 |         if len(y.shape) == 1:
101 |             squeeze = True
102 |             y = np.expand_dims(y, 0)
103 |         rv = np.zeros_like(y)
104 | 
105 |         for i in range(y.shape[0]):
106 |             # Assume in rad, rad/s:
107 |             #y = np.concatenate([np.broadcast_to(initial_positions, n), np.broadcast_to(initial_velocities, n)])
108 | 
109 |             vals = np.concatenate((y[i,:], parameter_vals))
110 |             sol = np.linalg.solve(mm_func(*vals), fo_func(*vals))
111 |             rv[i,:] = np.array(sol).T[0]
112 | 
113 |         if squeeze:
114 |             return rv[0,...]
115 |         return rv
116 | 
117 |     # ODE integration
118 |     return gradient
119 | 
120 | def _redim(inp):
121 |     vec = np.array(inp)
122 |     # Wrap all dimensions:
123 |     n = vec.shape[1] // 2
124 |     assert vec.shape[1] == n*2
125 | 
126 |     # Get angular positions:
127 |     pos = vec[:,:n]
128 |     l = 100
129 | 
130 |     if np.any(pos < -np.pi):
131 |         # In multiples of 2pi
132 |         adj, _ = np.modf((pos[pos < -np.pi] + np.pi) / (2*np.pi))
133 |         # Scale it back
134 |         pos[pos < -np.pi] = (adj * 2*np.pi) + np.pi
135 |         assert not np.any(pos < -np.pi)
136 | 
137 |     if np.any(pos >= np.pi):
138 |         # In multiples of 2pi
139 |         adj, _ = np.modf((pos[pos >= np.pi] - np.pi) / (2*np.pi))
140 |         # Scale it back
141 |         pos[pos >= np.pi] = (adj * 2*np.pi) - np.pi
142 |         assert not np.any(pos >= np.pi)
143 | 
144 |     vec[:,:n] = pos
145 |     return vec
146 | 
147 | NUM_EXAMPLES = lambda n: 200 * n
148 | def build(props):
149 |     # Number of joints in the pendulum:
150 |     n = int(props["n"]) if "n" in props else 1
151 |     test = "test" if "test" in props else "train"
152 |     lowenergy =  "lowenergy" in props
153 | 
154 |     pen_gen = pendulum_gradient(n)
155 |     le_str = "-lowenergy" if lowenergy else ""
156 |     cache_path = CACHE / f"p-{n}{le_str}-{test}.npz"
157 |     if not cache_path.exists():
158 |         if lowenergy:
159 |             X = np.zeros((NUM_EXAMPLES(n), 2 * n))
160 |             # Pick values for displacement in range [-pi/4, pi/4] radians
161 |             X[:,:n] = (np.random.rand(X.shape[0], n).astype(np.float32) - 0.5) * np.pi/2
162 |         else:
163 |             # Pick values in range [-pi, pi] radians, radians/sec
164 |             X = (np.random.rand(NUM_EXAMPLES(n), n * 2).astype(np.float32) - 0.5) * 2 * np.pi
165 |         Y = pen_gen(X)
166 | 
167 |         cache_path.parent.mkdir(parents=True, exist_ok=True)
168 |         np.savez(cache_path, X=X, Y=Y)
169 |     else:
170 |         load = np.load(cache_path)
171 |         X = load["X"]
172 |         Y = load["Y"]
173 | 
174 |     rv = torch.utils.data.TensorDataset(torch.tensor(X), torch.tensor(Y))
175 |     rv._pendulum_gen = pen_gen
176 |     rv._n = n
177 |     rv._redim = _redim
178 |     return rv
179 | 
180 | if __name__ == "__main__":
181 |     # Continuous-time dynamics, learning function [x, x']_t -> [x', x'']_{t+1}
182 |     def f_true(x):
183 |         g_l = 9.81
184 |         b = 0.3
185 |         return np.array([x[:,1], -b*x[:,1] + g_l*np.sin(x[:,0]-np.pi)]).T
186 | 
187 |     gfunc = pendulum_gradient(1)
188 |     X = (np.random.rand(10000, 2).astype(np.float32) - 0.5) * 2 * np.pi # Pick values in range [-pi, pi] radians, radians/sec
189 | 
190 |     test = gfunc(X)
191 |     ref = f_true(X)
192 | 
193 |     differences = sorted(np.sum((test - ref) ** 2, axis=1))[-8:]
194 |     print(differences)
195 |     assert differences[-1] < 1e-8
196 | 


--------------------------------------------------------------------------------
/models/pendulum_energy.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | 
  3 | import torch
  4 | from sympy import Dummy, lambdify, srepr, symbols
  5 | from sympy.core.function import AppliedUndef
  6 | from sympy.core.sympify import sympify
  7 | from sympy.physics import mechanics
  8 | from sympy.physics.vector import Vector
  9 | from sympy.printing.printer import Printer
 10 | 
 11 | try:
 12 |     import sympy2torch
 13 | except ImportError:
 14 |     from . import sympy2torch
 15 | 
 16 | import logging
 17 | logger = logging.getLogger(__name__)
 18 | 
 19 | def pendulum_energy(n=1, lengths=1, masses=1, include_gpe=True, include_ke=True):
 20 |     # Generalized coordinates and velocities
 21 |     # (in this case, angular positions & velocities of each mass) 
 22 |     q = mechanics.dynamicsymbols('q:{0}'.format(n))
 23 |     u = mechanics.dynamicsymbols('u:{0}'.format(n))
 24 | 
 25 |     # mass and length
 26 |     m = symbols('m:{0}'.format(n))
 27 |     l = symbols('l:{0}'.format(n))
 28 | 
 29 |     # gravity and time symbols
 30 |     g, t = symbols('g,t')
 31 | 
 32 |     #--------------------------------------------------
 33 |     # Step 2: build the model using Kane's Method
 34 | 
 35 |     # Create pivot point reference frame
 36 |     A = mechanics.ReferenceFrame('A')
 37 |     P = mechanics.Point('P')
 38 |     Origin = P
 39 |     P.set_vel(A, 0)
 40 | 
 41 |     gravity_direction = -A.x
 42 | 
 43 |     # lists to hold particles, forces, and kinetic ODEs
 44 |     # for each pendulum in the chain
 45 |     particles = []
 46 |     forces = []
 47 | 
 48 |     gpe = []
 49 |     ke = []
 50 | 
 51 |     cartVel = 0.0
 52 |     cartPos = 0.0
 53 | 
 54 |     for i in range(n):
 55 |         # Create a reference frame following the i^th mass
 56 |         Ai = A.orientnew('A' + str(i), 'Axis', [q[i], A.z])
 57 |         Ai.set_ang_vel(A, u[i] * A.z)
 58 | 
 59 |         # Create a point in this reference frame
 60 |         Pi = P.locatenew('P' + str(i), l[i] * Ai.x)
 61 |         Pi.v2pt_theory(P, A, Ai)
 62 | 
 63 |         # Create a new particle of mass m[i] at this point
 64 |         Pai = mechanics.Particle('Pa' + str(i), Pi, m[i])
 65 |         particles.append(Pai)
 66 | 
 67 |         # Calculate the cartesian position and velocity:
 68 |         # cartPos += l[i] * q[i]
 69 |         pos = Pi.pos_from(Origin)
 70 | 
 71 |         ke.append(1/n * Pai.kinetic_energy(A))
 72 |         gpe.append(m[i] * g * (Pi.pos_from(Origin) & gravity_direction))
 73 | 
 74 |         P = Pi
 75 | 
 76 | 
 77 |     # lengths and masses
 78 |     if lengths is None:
 79 |         lengths = np.ones(n) / n
 80 |     lengths = np.broadcast_to(lengths, n)
 81 |     masses = np.broadcast_to(masses, n)
 82 | 
 83 |     # Fixed parameters: gravitational constant, lengths, and masses
 84 |     parameters = [g] + list(l) + list(m)
 85 |     parameter_vals = [9.81] + list(lengths) + list(masses)
 86 | 
 87 |     # define symbols for unknown parameters
 88 |     unknowns = [Dummy() for i in q + u]
 89 |     unknown_dict = dict(zip(q + u, unknowns))
 90 | 
 91 |     # create functions for numerical calculation
 92 |     total_energy = 0
 93 |     if include_gpe:
 94 |         total_energy += (sum(gpe)).subs(zip(parameters, parameter_vals))
 95 |     if include_ke:
 96 |         total_energy += (sum( ke)).subs(zip(parameters, parameter_vals))
 97 | 
 98 |     total_energy_func = sympy2torch.sympy2torch(total_energy)
 99 | 
100 |     minimum_energy = total_energy_func(**fixvalue(n, torch.tensor([[0.]*2*n]))).detach()
101 |     return lambda inp: (total_energy_func(**fixvalue(n, inp)) - minimum_energy.to(inp)).unsqueeze(1)
102 | 
103 | def fixvalue(n, value):
104 |     keys = [f"q{i}_t" for i in range(n)] + [f"u{i}_t" for i in range(n)]
105 |     rv = {}
106 |     for i in range(2*n):
107 |         if isinstance(value, list):
108 |             rv[keys[i]] = value[i]
109 |         else:
110 |             rv[keys[i]] = value[:,i]
111 |     return rv
112 | 
113 | if __name__ == "__main__":
114 |     n = 2
115 |     efunc = pendulum_energy(n)
116 | 
117 |     # Test this!
118 |     print(efunc(torch.tensor([[0., 0., 0., 0.]])))
119 |     print(efunc(torch.tensor([[ np.pi/2,  np.pi/2, 0., 0.]])))
120 |     print(efunc(torch.tensor([[-np.pi/2, -np.pi/2, 0., 0.]])))
121 |     print(efunc(torch.tensor([[ np.pi  ,  np.pi  , 0., 0.]])))
122 | 
123 |     print()
124 | 
125 |     print(efunc(torch.tensor([[0., 0.,  1.,  0.]])))
126 |     print(efunc(torch.tensor([[0., 0., -1.,  0.]])))
127 |     print(efunc(torch.tensor([[0., 0.,  1.,  1.]])))
128 |     print(efunc(torch.tensor([[0., 0., -1., -1.]])))
129 | 
130 |     print()
131 | 
132 |     print(efunc(torch.tensor([[ np.pi/2,       0, 0., 0.]])))
133 |     print(efunc(torch.tensor([[ np.pi/2,       0, .1, 0.]])))
134 |     print(efunc(torch.tensor([[       0, np.pi/2, 0., 0.]])))
135 |     print(efunc(torch.tensor([[ np.pi/96, np.pi/2, 0., 0.]])))
136 |     print(efunc(torch.tensor([[-np.pi/96, np.pi/2, 0., 0.]])))
137 | 


--------------------------------------------------------------------------------
/models/simple.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import logging
 4 | 
 5 | import torch
 6 | from torch import nn
 7 | 
 8 | logger = logging.getLogger(__name__)
 9 | 
10 | global SIZE_A, SIZE_B, model
11 | SIZE_A = 320
12 | SIZE_B = 320
13 | 
14 | model = None
15 | 
16 | loss_ = nn.MSELoss()
17 | loss = lambda Ypred, Yactual, X, **kw: loss_(Ypred, Yactual)
18 | 
19 | def loss_flatten(l):
20 |     return [l]
21 | 
22 | def loss_labels():
23 |     return ["loss"]
24 | 
25 | def summary(*a, **kw):
26 |     pass
27 | 
28 | def configure(props):
29 |     global SIZE_A, SIZE_B, model
30 |     if "a" in props:
31 |         SIZE_A = int(props["a"])
32 | 
33 |     if "b" in props:
34 |         SIZE_B = int(props["b"])
35 | 
36 |     model = nn.Sequential(
37 |         nn.Linear(SIZE_A, SIZE_B), nn.ReLU(),
38 |         nn.Linear(SIZE_B, SIZE_B), nn.ReLU(),
39 |         nn.Linear(SIZE_B, SIZE_A))
40 | 
41 |     logger.info(f"Set layer sizes to {SIZE_A} -> {SIZE_B} -> {SIZE_B} -> {SIZE_A}")
42 | 
43 | 


--------------------------------------------------------------------------------
/models/stabledynamics.py:
--------------------------------------------------------------------------------
  1 | import logging
  2 | import math
  3 | 
  4 | import torch
  5 | import torch.nn.functional as F
  6 | from torch import nn
  7 | from . import pendulum_energy
  8 | 
  9 | logger = logging.getLogger(__file__)
 10 | 
 11 | #
 12 | # From StableDynamics.ipynb; Zico Kolter
 13 | #
 14 | 
 15 | # You can use this to compensate for numeric error:
 16 | VERIFY = False
 17 | V_SCALE = 0.01
 18 | 
 19 | global V_WRAP, SCALE_FX
 20 | V_WRAP = False
 21 | SCALE_FX = False
 22 | 
 23 | class Dynamics(nn.Module):
 24 |     def __init__(self, fhat, V, alpha=0.01):
 25 |         super().__init__()
 26 |         self.fhat = fhat
 27 |         self.V = V
 28 |         self.alpha = alpha
 29 | 
 30 |     def forward(self, x):
 31 |         fx = self.fhat(x)
 32 |         if SCALE_FX:
 33 |             fx = fx / fx.norm(p=2, dim=1, keepdim=True).clamp(min=1.0)
 34 | 
 35 |         Vx = self.V(x)
 36 |         gV = torch.autograd.grad([a for a in Vx], [x], create_graph=True, only_inputs=True)[0]
 37 |         rv = fx - gV * (F.relu((gV*fx).sum(dim=1) + self.alpha*Vx[:,0])/(gV**2).sum(dim=1))[:,None]
 38 | 
 39 |         if VERIFY:
 40 |             # Verify that rv has no positive component along gV.
 41 |             # This helps us catch:
 42 |             #   (1) numeric error in the symbolic gradient calculation, and
 43 |             #   (2) Violation of the Lyapunov function when Euler integration is used.
 44 |             verify = (gV * rv).sum(dim=1)
 45 |             num_violation = len([v for v in verify if v > 0]) # (1)
 46 |             new_V = self.V(x + V_SCALE * rv)
 47 |             if (new_V > Vx).any(): # (2)
 48 |                 err = sorted([v for v in (new_V - Vx).detach().cpu().numpy().ravel() if v > 0], reverse=True)
 49 |                 logger.warn(f"V increased by: {err[:min(5, len(err))]} (total {len(err)}; upward grad {num_violation});")
 50 | 
 51 |         return rv
 52 | 
 53 | class ICNN(nn.Module):
 54 |     def __init__(self, layer_sizes, activation=F.relu_):
 55 |         super().__init__()
 56 |         self.W = nn.ParameterList([nn.Parameter(torch.Tensor(l, layer_sizes[0])) 
 57 |                                    for l in layer_sizes[1:]])
 58 |         self.U = nn.ParameterList([nn.Parameter(torch.Tensor(layer_sizes[i+1], layer_sizes[i]))
 59 |                                    for i in range(1,len(layer_sizes)-1)])
 60 |         self.bias = nn.ParameterList([nn.Parameter(torch.Tensor(l)) for l in layer_sizes[1:]])
 61 |         self.act = activation
 62 |         self.reset_parameters()
 63 |         logger.info(f"Initialized ICNN with {self.act} activation")
 64 | 
 65 |     def reset_parameters(self):
 66 |         # copying from PyTorch Linear
 67 |         for W in self.W:
 68 |             nn.init.kaiming_uniform_(W, a=5**0.5)
 69 |         for U in self.U:
 70 |             nn.init.kaiming_uniform_(U, a=5**0.5)
 71 |         for i,b in enumerate(self.bias):
 72 |             fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.W[i])
 73 |             bound = 1 / (fan_in**0.5)
 74 |             nn.init.uniform_(b, -bound, bound)
 75 | 
 76 |     def forward(self, x):
 77 |         z = F.linear(x, self.W[0], self.bias[0])
 78 |         z = self.act(z)
 79 | 
 80 |         for W,b,U in zip(self.W[1:-1], self.bias[1:-1], self.U[:-1]):
 81 |             z = F.linear(x, W, b) + F.linear(z, F.softplus(U)) / U.shape[0]
 82 |             z = self.act(z)
 83 | 
 84 |         return F.linear(x, self.W[-1], self.bias[-1]) + F.linear(z, F.softplus(self.U[-1])) / self.U[-1].shape[0]
 85 | 
 86 | 
 87 | 
 88 | class ReHU(nn.Module):
 89 |     """ Rectified Huber unit"""
 90 |     def __init__(self, d):
 91 |         super().__init__()
 92 |         self.a = 1/d
 93 |         self.b = -d/2
 94 | 
 95 |     def forward(self, x):
 96 |         return torch.max(torch.clamp(torch.sign(x)*self.a/2*x**2,min=0,max=-self.b),x+self.b)
 97 | 
 98 | class MakePSD(nn.Module):
 99 |     def __init__(self, f, n, eps=0.01, d=1.0):
100 |         super().__init__()
101 |         self.f = f
102 |         self.zero = torch.nn.Parameter(f(torch.zeros(1,n)), requires_grad=False)
103 |         self.eps = eps
104 |         self.d = d
105 |         self.rehu = ReHU(self.d)
106 | 
107 |     def forward(self, x):
108 |         smoothed_output = self.rehu(self.f(x) - self.zero)
109 |         quadratic_under = self.eps*(x**2).sum(1,keepdim=True)
110 |         return smoothed_output + quadratic_under
111 | 
112 | class PosDefICNN(nn.Module):
113 |     def __init__(self, layer_sizes, eps=0.1, negative_slope=0.05):
114 |         super().__init__()
115 |         self.W = nn.ParameterList([nn.Parameter(torch.Tensor(l, layer_sizes[0])) 
116 |                                    for l in layer_sizes[1:]])
117 |         self.U = nn.ParameterList([nn.Parameter(torch.Tensor(layer_sizes[i+1], layer_sizes[i]))
118 |                                    for i in range(1,len(layer_sizes)-1)])
119 |         self.eps = eps
120 |         self.negative_slope = negative_slope
121 |         self.reset_parameters()
122 | 
123 |     def reset_parameters(self):
124 |         # copying from PyTorch Linear
125 |         for W in self.W:
126 |             nn.init.kaiming_uniform_(W, a=5**0.5)
127 |         for U in self.U:
128 |             nn.init.kaiming_uniform_(U, a=5**0.5)
129 | 
130 |     def forward(self, x):
131 |         z = F.linear(x, self.W[0])
132 |         F.leaky_relu_(z, negative_slope=self.negative_slope)
133 | 
134 |         for W,U in zip(self.W[1:-1], self.U[:-1]):
135 |             z = F.linear(x, W) + F.linear(z, F.softplus(U))*self.negative_slope
136 |             z = F.leaky_relu_(z, negative_slope=self.negative_slope)
137 | 
138 |         z = F.linear(x, self.W[-1]) + F.linear(z, F.softplus(self.U[-1]))
139 |         return F.relu(z) + self.eps*(x**2).sum(1)[:,None]
140 | 
141 | def loss(Ypred, Yactual, X):
142 |     # Force smoothness in V:
143 |     # penalty for new_V being larget than old V:
144 |     Vloss = torch.tensor(0)
145 |     if SMOOTH_V:
146 |         V = model.V
147 |         # Successor to X:
148 |         succ_X = (X + V_SCALE * Yactual).detach()
149 |         if V_WRAP:
150 |             while torch.any(succ_X < -math.pi):
151 |                 succ_X[succ_X < -math.pi] = succ_X[succ_X < -math.pi] + 2 * math.pi
152 |             while torch.any(succ_X >= math.pi):
153 |                 succ_X[succ_X >= math.pi] = succ_X[succ_X >= math.pi] - 2 * math.pi
154 |             succ_X.requires_grad_()
155 | 
156 |         Vloss = (V(succ_X) - V(X)).clamp(min=0).mean()
157 | 
158 |     l2loss = ((Ypred - Yactual)**2).mean()
159 | 
160 |     return (l2loss + SMOOTH_V * Vloss, l2loss, Vloss)
161 | 
162 | global model, SMOOTH_V
163 | model = None
164 | SMOOTH_V = 0
165 | 
166 | def loss_flatten(l):
167 |     return l
168 | 
169 | def loss_labels():
170 |     return ["loss", "l2", "V"]
171 | 
172 | def summary(*a, **kw):
173 |     pass
174 | 
175 | def configure(props):
176 |     logger.info(props)
177 |     lsd = int(props["latent_space_dim"])
178 |     logger.info(f"Set latent space dimenson to {lsd}")
179 | 
180 |     h_dim = int(props["h"]) if "h" in props else 100
181 |     ph_dim = int(props["hp"]) if "hp" in props else 40
182 |     logger.info(f"Set hidden layer size to {h_dim} and hidden layer in projection to {ph_dim}")
183 | 
184 |     if "scale_fx" in props and props["scale_fx"] not in ["false", "False"]:
185 |         logger.info(f"Scaling fx to prevent errors from too-large steps in Euler integration")
186 |         SCALE_FX = True
187 | 
188 |     # The function to learn
189 |     fhat = nn.Sequential(nn.Linear(lsd, h_dim), nn.ReLU(),
190 |                         nn.Linear(h_dim, h_dim), nn.ReLU(),
191 |                         nn.Linear(h_dim, lsd))
192 | 
193 |     ## The convex function to project onto:
194 |     projfn_eps = float(props["projfn_eps"]) if "projfn_eps" in props else 0.01
195 |     pendulum_n =   int(props["pendulum_n"]) if "pendulum_n" in props else None
196 |     if "projfn" in props:
197 |         if props["projfn"] == "PSICNN":
198 |             V = PosDefICNN([lsd, ph_dim, ph_dim, 1], eps=projfn_eps, negative_slope=0.3)
199 |         elif props["projfn"] == "ICNN":
200 |             V = ICNN([lsd, ph_dim, ph_dim, 1])
201 |         elif props["projfn"] == "PSD":
202 |             V = MakePSD(ICNN([lsd, ph_dim, ph_dim, 1]), lsd, eps=projfn_eps, d=1.0)
203 |         elif props["projfn"] == "PSD-REHU":
204 |             V = MakePSD(ICNN([lsd, ph_dim, ph_dim, 1], activation=ReHU(float(props["rehu"]) if "rehu" in props else 0.01)), lsd, eps=projfn_eps, d=1.0)
205 |         elif props["projfn"] == "NN-REHU":
206 |             seq = nn.Sequential(
207 |                     nn.Linear(lsd, ph_dim,), nn.ReLU(),
208 |                     nn.Linear(ph_dim, ph_dim), nn.ReLU(),
209 |                     nn.Linear(ph_dim, 1), ReHU(float(props["rehu"]) if "rehu" in props else 0.01))
210 |             V = MakePSD(seq, lsd, eps=projfn_eps, d=1.0)
211 | 
212 |         elif props["projfn"] == "EndPSICNN":
213 |             V = nn.Sequential(nn.Linear(lsd, ph_dim, bias=False), nn.LeakyReLU(),
214 |                 nn.Linear(ph_dim, lsd, bias=False), nn.LeakyReLU(),
215 |                 PosDefICNN([lsd, ph_dim, ph_dim, 1], eps=projfn_eps, negative_slope=0.3))
216 |         elif props["projfn"] == "NN":
217 |             V = nn.Sequential(
218 |                     nn.Linear(lsd, ph_dim,), nn.ReLU(),
219 |                     nn.Linear(ph_dim, ph_dim), nn.ReLU(),
220 |                     nn.Linear(ph_dim, 1))
221 |         elif props["projfn"] == "ActualPendulumEnergy":
222 |             V = pendulum_energy.pendulum_energy(n=pendulum_n)
223 |         else:
224 |             logger.error(f"Projected function {props['projfn']} does not exist")
225 | 
226 |     logger.info(f"Set Lyapunov function to {V} with param eps={projfn_eps}, n={pendulum_n}")
227 | 
228 |     alpha = float(props["a"]) if "a" in props else 0.01
229 |     logger.info(f"Set alpha to {alpha}")
230 | 
231 |     global SMOOTH_V, V_WRAP
232 |     if "wrap" in props:
233 |         V_WRAP = True
234 | 
235 |     if "smooth_v" in props:
236 |         SMOOTH_V = float(props["smooth_v"])
237 |         logger.info(f"Set smooth_v loss coeff to {SMOOTH_V}")
238 |         if V_WRAP:
239 |             logger.warning("V_WRAP IS SET; smoothing of V will wrap from -pi to pi")
240 | 
241 |     global model
242 |     model = Dynamics(fhat, V, alpha=alpha)
243 | 


--------------------------------------------------------------------------------
/models/sympy2torch.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | 
  3 | import torch
  4 | from sympy import Dummy, lambdify, srepr, symbols
  5 | from sympy.core.function import AppliedUndef
  6 | from sympy.core.sympify import sympify
  7 | from sympy.core.power import Pow
  8 | from sympy.physics import mechanics
  9 | from sympy.physics.vector import Vector
 10 | from sympy.printing.printer import Printer
 11 | 
 12 | class TorchPrinter(Printer):
 13 |     printmethod = "_torchrepr"
 14 |     _default_settings = {
 15 |         "order": None
 16 |     }
 17 | 
 18 |     def _print_Add(self, expr, order=None):
 19 |         terms = list(map(self._print, self._as_ordered_terms(expr, order=order)))
 20 |         def __inner_sum(**kw):
 21 |             x = terms[0](**kw)
 22 |             for r in terms[1:]:
 23 |                 v = r(**kw)
 24 |                 x = torch.add(x, v)
 25 |             return x
 26 |         return __inner_sum
 27 | 
 28 |     def _print_Function(self, expr):
 29 |         __FUNCTION_MAP = { "sin": torch.sin, "cos": torch.cos, "tan": torch.tan, "Abs": torch.abs }
 30 | 
 31 |         if expr.func.__name__ in __FUNCTION_MAP:
 32 |             func = __FUNCTION_MAP[expr.func.__name__]
 33 |             args = [self._print(a) for a in expr.args]
 34 |             return lambda **kw: func(*(a(**kw) for a in args))
 35 |         else:
 36 |             key = f"{expr.func.__name__}_{', '.join([str(a) for a in expr.args])}"
 37 |             return lambda **kw: kw[key]
 38 | 
 39 |     def _print_Half(self, expr):
 40 |         return lambda **kw: 0.5
 41 | 
 42 |     def _print_Integer(self, expr):
 43 |         return lambda **kw: expr.p
 44 | 
 45 |     def _print_NaN(self, expr):
 46 |         return lambda **kw: float('nan')
 47 | 
 48 |     def _print_Mul(self, expr, order=None):
 49 |         assert len(expr.args) > 1
 50 |         terms = [self._print(a) for a in expr.args]
 51 | 
 52 |         def __inner_mul(**kw):
 53 |             x = terms[0](**kw)
 54 |             for r in terms[1:]:
 55 |                 x = torch.mul(x, r(**kw))
 56 |             return x
 57 |         return __inner_mul
 58 | 
 59 |     def _print_Rational(self, expr):
 60 |         return lambda **kw: expr.p/expr.q
 61 | 
 62 |     def _print_Fraction(self, expr):
 63 |         numer = self._print(expr.numerator)
 64 |         denom = self._print(expr.denominator)
 65 |         return lambda **kw: numer(**kw), denom(**kw)
 66 | 
 67 |     def _print_Float(self, expr):
 68 |         return lambda **kw: float(expr.evalf())
 69 | 
 70 |     def _print_Symbol(self, expr):
 71 |         d = expr._assumptions.generator
 72 |         # print the dummy_index like it was an assumption
 73 |         if expr.is_Dummy or d != {}:
 74 |             raise NotImplementedError()
 75 | 
 76 |         return lambda **kw: kw[expr.name]
 77 | 
 78 |     def _print_FracElement(self, frac):
 79 |         numer_terms = list(frac.numer.terms())
 80 |         assert len(numer_terms) == 1
 81 |         #numer_terms.sort(key=frac.field.order, reverse=True)
 82 |         denom_terms = list(frac.denom.terms())
 83 |         assert len(denom_terms) == 1
 84 |         #denom_terms.sort(key=frac.field.order, reverse=True)
 85 |         numer = self._print(numer[0])
 86 |         denom = self._print(denom[0])
 87 |         return lambda *kw: torch.div(numer(*kw), denom(*kw))
 88 | 
 89 |     def _print_Expr(self, expr):
 90 |         base = self._print(expr.base)
 91 |         exp = self._print(expr.exp)
 92 |         if isinstance(expr, Pow):
 93 |             return lambda **kw: torch.pow(base(**kw), exp(**kw))
 94 | 
 95 |         raise NotImplementedError(f"No implementation of Expr {expr}")
 96 |         # return lambda **kw: kw[expr.name]
 97 | 
 98 | def sympy2torch(expr):
 99 |     return TorchPrinter()._print(expr)
100 | 
101 | def test_sympy2torch():
102 |     def test(expr, vars, examples):
103 |         func = sympy2torch(sympify(expr))
104 |         for i, (inp, out) in enumerate(examples):
105 |             kw = dict(zip(vars, torch.tensor(inp)))
106 |             v = func(**kw)
107 |             if v != out:
108 |                 print(f"FAILED\t{expr}@{i}\t{kw}\t Expected({out}) != Received({v})")
109 | 
110 |     test("Mul(Integer(2), Symbol('x'))", "x", [([0], 0), ([1], 2), ([2], 4)])
111 |     test("Add(Integer(2), Symbol('y'))", "xy", [([0, 0], 2), ([1, 0], 2), ([2, 2], 4), ([2, 1], 3)])
112 |     test("Mul(Add(Integer(2), Symbol('y')), Symbol('x'))", "xy", [([0, 0], 0), ([1, 0], 2), ([2, 0], 4), ([2, 1], 6)])
113 |     test("Mul(Integer(-1), Float('9.8100000000000005', precision=53), sin(Function('q0')(Symbol('t'))))", ["q0_t"], [([0.], 0), ([np.pi/2], -9.81), ([np.pi/6], -9.81*0.5), ([np.pi/4], -9.81*(0.5**0.5))])
114 |     test("Symbol('t')**2", "t", [([0.], 0), ([-2.], 4), ([2.], 4)])
115 | 
116 |     print("Completed!")
117 | 
118 | if __name__ == "__main__":
119 |     test_sympy2torch()
120 | 


--------------------------------------------------------------------------------
/models/train_many:
--------------------------------------------------------------------------------
 1 | ./train_pendulum_stable_knownV knownv 1 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05
 2 | ./train_pendulum_stable_knownV knownv 2 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05
 3 | ./train_pendulum_stable_knownV knownv 4 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05
 4 | ./train_pendulum_stable_knownV knownv 6 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05
 5 | 
 6 | ./train_pendulum_simple 1
 7 | ./train_pendulum_simple 2
 8 | ./train_pendulum_simple 4
 9 | ./train_pendulum_simple 6
10 | ./train_pendulum_simple 8
11 | 
12 | ./train_pendulum_stable stable 1 0.001 PSD-REHU 0.001 0 0.005
13 | ./train_pendulum_stable stable 2 0.001 PSD-REHU 0.001 0 0.005
14 | ./train_pendulum_stable stable 4 0.001 PSD-REHU 0.001 0 0.005
15 | ./train_pendulum_stable stable 6 0.001 PSD-REHU 0.001 0 0.005
16 | ./train_pendulum_stable stable 8 0.001 PSD-REHU 0.001 0 0.005
17 | 


--------------------------------------------------------------------------------
/models/vae.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | 
  3 | import logging
  4 | 
  5 | import torch
  6 | import torch.nn.functional as F
  7 | from torch import nn
  8 | from torch.autograd import Variable
  9 | 
 10 | logger = logging.getLogger(__name__)
 11 | 
 12 | class VAE(nn.Module):
 13 |     def __init__(self, LATENT_SPACE_DIM=320):
 14 |         super().__init__()
 15 | 
 16 |         self.fc_e1 = nn.Conv2d(3, 8, 9, stride=2)
 17 |         self.fc_e2 = nn.Conv2d(8, 16, 9, stride=2)
 18 |         self.fc_e3 = nn.Conv2d(16, 32, 5, stride=2)
 19 |         self.fc_e4 = nn.Conv2d(32, 64, 5, stride=2)
 20 |         self.fc_e5 = nn.Conv2d(64, 128, 3, stride=1)
 21 |         self.fc_e61 = nn.Linear(9*14*128, LATENT_SPACE_DIM)
 22 |         self.fc_e62 = nn.Linear(9*14*128, LATENT_SPACE_DIM)
 23 | 
 24 |         self.fc_d1 = nn.Linear(LATENT_SPACE_DIM, 9*14*128)
 25 |         self.fc_d2 = nn.ConvTranspose2d(128, 64, 3, stride=1)
 26 |         self.fc_d3 = nn.ConvTranspose2d(64, 32, 5, stride=2)
 27 |         self.fc_d4 = nn.ConvTranspose2d(32, 16, 5, stride=2)
 28 |         self.fc_d5 = nn.ConvTranspose2d(16,  8, 9, stride=2)
 29 |         self.fc_d6 = nn.ConvTranspose2d( 8,  3, 9, stride=2)
 30 | 
 31 |     def encode(self, x):
 32 |         assert list(x.size())[1:] == [3, 240, 320]
 33 |         x = F.relu(self.fc_e1(x))
 34 |         x = F.relu(self.fc_e2(x))
 35 |         x = F.relu(self.fc_e3(x))
 36 |         x = F.relu(self.fc_e4(x))
 37 |         x = F.relu(self.fc_e5(x))
 38 |         x = x.view([x.size()[0], -1])
 39 |         return self.fc_e61(x), self.fc_e62(x)
 40 | 
 41 |     def reparametrize(self, mu, logvar):
 42 |         std = logvar.mul(0.5).exp_()
 43 |         if torch.cuda.is_available():
 44 |             eps = torch.cuda.FloatTensor(std.size()).normal_()
 45 |         else:
 46 |             eps = torch.FloatTensor(std.size()).normal_()
 47 |         eps = Variable(eps)
 48 |         return eps.mul(std).add_(mu)
 49 | 
 50 |     def decode(self, z):
 51 |         nb = z.size()[0]
 52 |         z = self.fc_d1(z)
 53 |         z = z.view([nb, 128, 9, 14])
 54 |         z = F.relu(self.fc_d2(z, output_size=[nb, 32, 11, 16]))
 55 |         z = F.relu(self.fc_d3(z, output_size=[nb, 32, 25, 35]))
 56 |         z = F.relu(self.fc_d4(z, output_size=[nb, 16, 54, 74]))
 57 |         z = F.relu(self.fc_d5(z, output_size=[nb,  8, 116, 156]))
 58 |         z = F.relu(self.fc_d6(z, output_size=[nb,  3, 240, 320]))
 59 |         z = torch.sigmoid(z)
 60 |         return z
 61 | 
 62 |     def forward(self, x):
 63 |         mu, logvar = self.encode(x)
 64 |         z = self.reparametrize(mu, logvar)
 65 |         return self.decode(z), mu, logvar, z
 66 | 
 67 | # model is a torch.nn.Module that contains the model definition.
 68 | global model, VAE_BCE_LOSS
 69 | model = VAE()
 70 | VAE_BCE_LOSS = 1.
 71 | 
 72 | # Use MSE loss as distance from input to output:
 73 | reconstruction_function = lambda recon_x, x: torch.sum((recon_x - x) ** 2)
 74 | def vae_loss_fn(recon_x, x, mu, logvar):
 75 |     """
 76 |     recon_x: generating images
 77 |     x: origin images
 78 |     mu: latent mean
 79 |     logvar: latent log variance
 80 |     """
 81 |     BCE = torch.sum(reconstruction_function(recon_x, x))
 82 |     KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
 83 |     KLD = torch.sum(KLD_element).mul_(-0.5)
 84 |     # KL divergence
 85 |     return VAE_BCE_LOSS * BCE + KLD, BCE, KLD
 86 | 
 87 | def loss(Ypred, Yactual, X):
 88 |     """loss function for learning problem
 89 | 
 90 |     Arguments:
 91 |         Ypred {Model output type} -- predicted output
 92 |         Yactual {Data output type} -- output from data
 93 |         X {torch.Variable[input]} -- input
 94 | 
 95 |     Returns:
 96 |         Tuple[nn.Variable] -- Parts of the loss function; the first element is passed to the optimizer
 97 |         nn.Variable -- the loss to optimize
 98 |     """
 99 |     reconstructed, mu, logvar, z = Ypred
100 |     return vae_loss_fn(reconstructed, Yactual, mu, logvar)
101 | 
102 | def loss_labels():
103 |     """Keys corresponding to the components of the loss
104 |     
105 |     Returns:
106 |         Tuple[str] -- Tuple of loss 
107 |     """
108 |     return ("vae_loss", "reconstruction_loss", "prior_loss")
109 | 
110 | def configure(props):
111 |     global model
112 | 
113 |     try:
114 |         lsd = props["latent_space_dim"]
115 |         model = VAE(lsd)
116 |     except KeyError:
117 |         pass
118 |     logger.info(f"Latent space is 1x{lsd}.")
119 | 
120 |     try:
121 |         global VAE_BCE_LOSS
122 |         VAE_BCE_LOSS = 10**float(props["bce"])
123 |     except KeyError:
124 |         pass
125 |     logger.info(f"BCE to KLD ratio {VAE_BCE_LOSS}.")
126 | 


--------------------------------------------------------------------------------
/models/vae_trajectory.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import logging
 4 | 
 5 | import torch
 6 | from torch import nn
 7 | 
 8 | from . import simple as dynmod_simple
 9 | from . import stabledynamics as dynmod_stable
10 | from . import vae as vaemod
11 | 
12 | logger = logging.getLogger(__name__)
13 | 
14 | class TrajectoryVAE(nn.Module):
15 |     def __init__(self, vae, dyn):
16 |         super().__init__()
17 |         self.vae = vae
18 |         self.dyn = dyn
19 | 
20 |     def forward(self, X):
21 |         X_a, X_b = X
22 |         Y_a, mu_a, logvar_a, z_a = self.vae(X_a)
23 |         z_b = z_a + self.dyn(z_a)
24 |         Y_b = self.vae.decode(z_b)
25 | 
26 |         return (Y_a, mu_a, logvar_a, z_a, Y_b, z_b)
27 | 
28 | global model, WEIGHT_NEXT
29 | model = None
30 | WEIGHT_NEXT = None
31 | 
32 | def loss(Ypred, Yactual, X):
33 |     Y_a, mu_a, logvar_a, z_a, Y_b, z_b = Ypred
34 |     X_a, X_b = X
35 |     YA_a, YA_b = Yactual
36 | 
37 |     ls_a, ls_a_bce, ls_a_kld = vaemod.loss((Y_a, mu_a, logvar_a, z_a), X_a, X_a)
38 |     ls_b_bce = vaemod.reconstruction_function(Y_b, YA_b)
39 | 
40 |     return (ls_a + WEIGHT_NEXT * ls_b_bce, ls_a, ls_a_bce, ls_a_kld, ls_b_bce, torch.sum((z_a - z_b)**2))
41 | 
42 | def summary(epoch, summarywriter, Ypred, X):
43 |     Y_a, mu_a, logvar_a, z_a, Y_b, z_b = Ypred
44 |     X_now, X_next = X
45 |     summarywriter.add_embedding(z_a.data, label_img=X_now.data, global_step=epoch, tag="learned_embedding")
46 |     summarywriter.add_images("current_reconstructed", Y_a.clamp(max=1.0), global_step=epoch)
47 |     summarywriter.add_images("next_reconstructed", Y_b.clamp(max=1.0), global_step=epoch)
48 | 
49 | def loss_flatten(l):
50 |     return l
51 | 
52 | def loss_labels():
53 |     return ["traj", "vae_a", "vae_a_bce", "vae_a_kld", "vae_b_bce", "z_dist"]
54 | 
55 | def configure(props):
56 |     dynmod = dynmod_stable if "stable" in props else dynmod_simple
57 |     logger.info(f"Latent space dynamics {dynmod.__file__}")
58 |     logger.info(f"Latent space dynamics {vaemod.__file__}")
59 | 
60 |     lsd = int(props["latent_space_dim"]) if "latent_space_dim" in props else 320
61 |     logger.info(f"Set latent space dim to {lsd}")
62 |     vaemod.configure({ **props, "latent_space_dim": lsd })
63 |     dynmod.configure({ **props, "latent_space_dim": lsd })
64 | 
65 |     if "w" in props:
66 |         global model, WEIGHT_NEXT
67 |         WEIGHT_NEXT = 10**float(props["w"])
68 |         logger.info(f"Set weight of next step reconstruction to {WEIGHT_NEXT}")
69 |     else:
70 |         logger.warn("No weight of next step reconstruction set; loss function will not run")
71 | 
72 |     model = TrajectoryVAE(vaemod.model, dynmod.model)
73 | 


--------------------------------------------------------------------------------
/pendulum_error.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | 
  3 | import argparse
  4 | import os
  5 | 
  6 | import matplotlib.cm as cm
  7 | import matplotlib.pyplot as plt
  8 | import numpy as np
  9 | import torch
 10 | from plot_data import plot_data, plot_data_args
 11 | from torchvision.utils import save_image
 12 | from util import (DynamicLoad, latest_file, loadDataFile, setup_logging, to_variable)
 13 | from scipy.integrate import odeint
 14 | 
 15 | from models import pendulum_energy
 16 | 
 17 | from pathlib import Path
 18 | 
 19 | logger = setup_logging(os.path.basename(__file__))
 20 | 
 21 | 
 22 | def main(args):
 23 |     model = args.model.model
 24 |     model.load_state_dict(torch.load(args.weight))
 25 |     model.eval()
 26 |     if torch.cuda.is_available():
 27 |         model.cuda()
 28 | 
 29 |     physics = args.data._pendulum_gen
 30 |     n = args.data._n
 31 |     redim = args.data._redim
 32 |     h = args.timestep
 33 | 
 34 |     logger.info(f"Loaded physics simulator for {n}-link pendulum")
 35 | 
 36 |     cache_path = Path("pendulum-cache") / f"p-physics-{n}.npy"
 37 | 
 38 |     # Energy functions
 39 |     energy = pendulum_energy.pendulum_energy(n)
 40 | 
 41 |     if not cache_path.exists():
 42 |         logger.info(f"Generating trajectories for {cache_path}")
 43 |         # Initialize args.number initial positions:
 44 |         X_init = np.zeros((args.number, 2 * n)).astype(np.float32)
 45 |         X_init[:,:] = (np.random.rand(args.number, 2*n).astype(np.float32) - 0.5) * np.pi/4 # Pick values in range [-pi/8, pi/8] radians, radians/sec
 46 | 
 47 |         X_phy = np.zeros((args.steps, *X_init.shape), dtype=np.float32)
 48 |         X_phy[0,...] = X_init
 49 |         for i in range(1, args.steps):
 50 |             logger.info(f"Timestep {i}")
 51 |             k1 = h * physics(X_phy[i-1,...])
 52 |             k2 = h * physics(X_phy[i-1,...] + k1/2)
 53 |             k3 = h * physics(X_phy[i-1,...] + k2/2)
 54 |             k4 = h * physics(X_phy[i-1,...] + k3)
 55 |             X_phy[i,...] = X_phy[i-1,...] + 1/6*(k1 + 2*k2 + 2*k3 + k4)
 56 |             assert not np.any(np.isnan(X_phy[i,...]))
 57 | 
 58 |         np.save(cache_path, X_phy)
 59 |         logger.info(f"Done generating trajectories for {cache_path}")
 60 | 
 61 |     else:
 62 |         X_phy = np.load(cache_path).astype(np.float32)
 63 |         logger.info(f"Loaded trajectories from {cache_path}")
 64 | 
 65 |     X_nn = to_variable(torch.tensor(X_phy[0,:,:]), cuda=torch.cuda.is_available())
 66 |     errors = np.zeros((args.steps,))
 67 |     for i in range(1, args.steps):
 68 |         X_nn.requires_grad = True
 69 |         k1 = h * model(X_nn)
 70 |         k1 = k1.detach()
 71 |         k2 = h * model(X_nn + k1/2)
 72 |         k2 = k2.detach()
 73 |         k3 = h * model(X_nn + k2/2)
 74 |         k3 = k3.detach()
 75 |         k4 = h * model(X_nn + k3)
 76 |         k4 = k4.detach()
 77 |         X_nn = X_nn + 1/6*(k1 + 2*k2 + 2*k3 + k4)
 78 |         X_nn = X_nn.detach()
 79 | 
 80 |         logger.info(f"Timestep {i}")
 81 |         y = X_nn.cpu().numpy()
 82 |         # TODO: Update error calculation
 83 |         vel_error = np.sum((X_phy[i,:,n:] - y[:,n:])**2)
 84 |         ang_error = (X_phy[i,:,:n] - y[:,:n])
 85 |         while np.any(ang_error >= np.pi):
 86 |             ang_error[ang_error >= np.pi] -= 2*np.pi
 87 |         while np.any(ang_error < -np.pi):
 88 |             ang_error[ang_error < -np.pi] += 2*np.pi
 89 | 
 90 |         ang_error = np.sum(ang_error**2)
 91 |         errors[i] = (vel_error + ang_error)
 92 | 
 93 |     for i in range(args.steps):
 94 |         print(f"{i}\t{np.sum(errors[0:i])}\t{errors[i]}")
 95 | 
 96 | 
 97 | if __name__ == "__main__":
 98 |     parser = argparse.ArgumentParser(description='Error of .')
 99 |     parser.add_argument('--number', type=int, default=1000, help="number of starting positions to evaluate from")
100 |     parser.add_argument('--timestep', type=float, default=0.01, help="duration of each timestep")
101 | 
102 |     parser.add_argument('data', type=DynamicLoad("datasets"), help='the pendulum dataset to load the simulator from')
103 |     parser.add_argument('model', type=DynamicLoad("models"), help='model to load')
104 |     parser.add_argument('weight', type=latest_file, help='model weight to load')
105 |     parser.add_argument('steps', type=int, help="number of steps to evaluate over")
106 | 
107 |     main(parser.parse_args())
108 | 


--------------------------------------------------------------------------------
/plot_data.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import argparse
 4 | 
 5 | import matplotlib.cm as cm
 6 | import matplotlib.pyplot as plt
 7 | import numpy as np
 8 | from sklearn.decomposition import PCA
 9 | from sklearn.manifold import TSNE
10 | 
11 | from util import loadDataFile
12 | import h5py
13 | 
14 | def main(args):
15 |     seq, param = args.data
16 |     plot_data(args, seq[args.select,:,:])
17 | 
18 | def plot_data(args, seq, original_seq = None):
19 |     if args.delta:
20 |         plt.figure()
21 |         dist = np.sum((seq[1:,...] - seq[:-1,...]) ** 2, axis=1) ** 0.5
22 |         plt.hist(dist)
23 |         if args.save:
24 |             plt.savefig(args.save, dpi=120)
25 |         else:
26 |             plt.show()
27 |         return
28 | 
29 |     if args.pca:
30 |         pca = PCA(n_components=args.pca)
31 |         if original_seq is None:
32 |             seq = pca.fit_transform(seq)
33 |         else:
34 |             original_seq = pca.fit_transform(original_seq)
35 |             seq = pca.transform(seq)
36 | 
37 |     if args.tsne:
38 |         tsne = TSNE(n_components=2, perplexity=30.0, n_iter=2000, verbose=2)
39 |         if original_seq is None:
40 |             seq = tsne.fit_transform(seq)
41 |         else:
42 |             tsne.fit(original_seq)
43 |             seq = tsne.transform(seq)
44 | 
45 |     if seq.shape[1] == 2:
46 |         plt.figure()
47 |         x, y = zip(*seq[:,:])
48 |         color_list = cm.get_cmap(name="viridis")
49 |         if args.strip:
50 |             n, m = tuple(args.strip)
51 |             for i in range(0, seq.shape[0] - 1, m):
52 |                 plt.plot(x[i:(i+n)], y[i:(i+n)], '-', color=color_list(i/(seq.shape[0]-1)))
53 |         else:
54 |             for i in range(seq.shape[0] - 1):
55 |                 plt.plot(x[i:(i+2)], y[i:(i+2)], '.', color=color_list(i/(seq.shape[0]-1)))
56 |         plt.axis('equal')
57 |         if args.save:
58 |             plt.savefig(args.save, dpi=120)
59 |         else:
60 |             plt.show()
61 |     else:
62 |         print("Cannot plot sequence: data is of size {}".format(seq.shape))
63 | 
64 | def plot_data_args(parser):
65 |     parser.add_argument('--pca', type=int, help="preprocess trajectories down to these many dimensions using PCA; happens before TSNE")
66 |     parser.add_argument('--tsne', action="store_true", help="use tsne to preprocess trajectories down to two dimensions")
67 |     parser.add_argument('--save', type=str, help="save image")
68 |     parser.add_argument('--strip', type=int, nargs=2, help="draw the first n elements in blocks of size m")
69 |     parser.add_argument('--delta', action="store_true", help="provide a histogram of the distribution of distances between points along the trajectory")
70 |     return parser
71 | 
72 | if __name__ == "__main__":
73 |     parser = plot_data_args(argparse.ArgumentParser(description='Plot data files.'))
74 |     parser.add_argument('data', type=loadDataFile, help='the data file to load')
75 |     parser.add_argument('--select', type=int, default=0, help="draw the trajectory at this index")
76 |     
77 |     main(parser.parse_args())
78 | 


--------------------------------------------------------------------------------
/render.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import argparse
 4 | import datetime
 5 | import glob
 6 | import os
 7 | import random
 8 | import sys
 9 | from collections import defaultdict
10 | from pathlib import Path
11 | 
12 | import numpy as np
13 | import torch
14 | import torch.nn.functional as F
15 | from torch import nn, optim
16 | from torch.autograd import Variable
17 | from torch.utils.data import DataLoader
18 | from torchvision.utils import save_image
19 | from util import DynamicLoad, setup_logging, to_variable, latest_file
20 | 
21 | logger = setup_logging(os.path.basename(__file__))
22 | 
23 | def main(args):
24 |     model = args.model.model
25 |     model.load_state_dict(torch.load(args.weight))
26 |     model.eval()
27 | 
28 |     if torch.cuda.is_available():
29 |         model.cuda()
30 | 
31 |     vaemod = model.vae
32 | 
33 |     if args.random is not None:
34 |         X = torch.tensor(np.random.normal(loc=0.0, scale=args.random, size=[args.batch_size, 320]).astype(np.float32))
35 |         X = to_variable(X, torch.cuda.is_available())
36 |         output = vaemod.decode(X)
37 | 
38 |     else:
39 |         test_dataloader = DataLoader(args.dataset, batch_size=args.batch_size, shuffle=False)
40 | 
41 |         for batch_idx, data in enumerate(test_dataloader):
42 |             X, Yactual = data
43 |             X = to_variable(X, torch.cuda.is_available())
44 |             output = vaemod(X)
45 |             break
46 | 
47 |     if isinstance(output, tuple):
48 |         output = output[0]
49 |     save_image(output, args.save)
50 | 
51 | 
52 | if __name__ == "__main__":
53 |     parser = argparse.ArgumentParser(description='Render the images passed through the model')
54 |     parser.add_argument('dataset', type=DynamicLoad("datasets"), help='dataset to train on')
55 |     parser.add_argument('model', type=DynamicLoad("models"), help='model to train with')
56 |     parser.add_argument('weight', type=latest_file, help='save model weight')
57 |     parser.add_argument('--batch-size', type=int, default=256, help='batch size')
58 |     parser.add_argument('--save', type=str, help="save renders as a single image")
59 |     parser.add_argument('--random', type=float, help="sample the space randomly, with zero mean and provided variance")
60 | 
61 |     try:
62 |         args = parser.parse_args()
63 |         main(args)
64 |     except:
65 |         raise
66 | 


--------------------------------------------------------------------------------
/render_vae:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env sh
 2 | 
 3 | set -e
 4 | 
 5 | DATASET="$1"
 6 | shift
 7 | MODEL="$1"
 8 | shift
 9 | 
10 | for INDIR in "$@";
11 | do
12 |     echo "PROCESSING $INDIR"
13 | 
14 |     ./.colorize ./trajectory.py                 \
15 |         --batch-size 360                        \
16 |         imagepairs[files=youtube/$DATASET/*]    \
17 |         "$MODEL"                                \
18 |         "$INDIR/checkpoint_*.pth"               \
19 |         "$INDIR/trajectory.h5"
20 | 
21 |     ./.colorize ./plot_data.py              \
22 |         --pca 2                             \
23 |         --save "$INDIR/plot.pca.png"        \
24 |         "$INDIR/trajectory.h5"
25 | 
26 |     ./.colorize ./plot_data.py              \
27 |         --pca 2                             \
28 |         --strip 50 500                      \
29 |         --save "$INDIR/plot.strip.png"      \
30 |         "$INDIR/trajectory.h5"
31 | 
32 |     ./.colorize ./vae_trajmodel.py              \
33 |         "$INDIR/trajectory.h5"                  \
34 |         "$MODEL"                                \
35 |         "$INDIR/checkpoint_*.pth"               \
36 |         plot                                    \
37 |         --pca 2                                 \
38 |         --save "$INDIR/plot.trajmodel.png"
39 | 
40 |     ./.colorize ./vae_trajmodel.py              \
41 |         "$INDIR/trajectory.h5"                  \
42 |         "$MODEL"                                \
43 |         "$INDIR/checkpoint_*.pth"               \
44 |         render                                  \
45 |         --save "$INDIR/trajmodel_render.png"
46 | 
47 |     ./.colorize ./render.py                     \
48 |         imagepairs[files=youtube/$DATASET/*]    \
49 |         "$MODEL"                                \
50 |         "$INDIR/checkpoint_*.pth"               \
51 |         --save "$INDIR/vae_render.png"
52 | 
53 |     ./.colorize ./render.py                     \
54 |         imagepairs[files=youtube/$DATASET/*]    \
55 |         "$MODEL"                                \
56 |         "$INDIR/checkpoint_*.pth"               \
57 |         --random 3                              \
58 |         --save "$INDIR/vae_render_random.png"
59 | 
60 |     mkdir -p "$INDIR/frames"
61 |     ./.colorize ./vae_trajmodel.py              \
62 |         --steps 300                             \
63 |         "$INDIR/trajectory.h5"                  \
64 |         "$MODEL"                                \
65 |         "$INDIR/checkpoint_*.pth"               \
66 |         render                                  \
67 |         --save-frames "$INDIR/frames/fr_{:0>05}.png"
68 |     ffmpeg -r 30 -f image2 -i "$INDIR/frames/fr_%05d.png" "$INDIR/generated.mp4"
69 | 
70 | done
71 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | matplotlib
2 | numpy
3 | torch
4 | torchvision
5 | tensorflow
6 | tensorboardX
7 | scikit-learn
8 | sympy
9 | 


--------------------------------------------------------------------------------
/test_lstm.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import argparse
 4 | import os
 5 | from pathlib import Path
 6 | 
 7 | import matplotlib.cm as cm
 8 | import matplotlib.pyplot as plt
 9 | import numpy as np
10 | import torch
11 | from models import pendulum_energy
12 | from plot_data import plot_data, plot_data_args
13 | from scipy.integrate import odeint
14 | from torchvision.utils import save_image
15 | from train_lstm import STEPS, TRAJECTORIES, LSTMModel, N
16 | from util import (DynamicLoad, latest_file, loadDataFile, setup_logging,
17 |                   to_variable)
18 | 
19 | logger = setup_logging(os.path.basename(__file__))
20 | 
21 | 
22 | def main(args):
23 |     hidden_dim = 116
24 |     model = LSTMModel(2*N, hidden_dim)
25 | 
26 |     model.load_state_dict(torch.load(args.weight))
27 |     model.eval()
28 |     if torch.cuda.is_available():
29 |         model.cuda()
30 | 
31 |     physics = args.data._pendulum_gen
32 |     n = args.data._n
33 |     redim = args.data._redim
34 |     h = args.timestep
35 | 
36 |     logger.info(f"Loaded physics simulator for {n}-link pendulum")
37 | 
38 |     cache_path = Path("pendulum-cache") / f"p-physics-{n}.npy"
39 | 
40 |     # Energy functions
41 |     energy = pendulum_energy.pendulum_energy(n)
42 | 
43 |     if cache_path.exists():
44 |         X_phy = np.load(cache_path).astype(np.float32)
45 |         logger.info(f"Loaded trajectories from {cache_path}")
46 |     else:
47 |         raise Exception(f"No trajectories for {cache_path}")
48 | 
49 |     X_nn = to_variable(torch.tensor(X_phy[0,:,:]), cuda=torch.cuda.is_available())
50 |     errors = np.zeros((args.steps,))
51 |     X_nn.requires_grad = True
52 |     X_nn = X_nn.unsqueeze(0)
53 | 
54 |     hiddens = None
55 |     for i in range(1, args.steps):
56 |         k1, new_hiddens = model(X_nn, hiddens)
57 |         k1 = h * k1.detach()
58 |         k2, _ = model(X_nn + k1/2, hiddens)
59 |         k2 = h * k2.detach()
60 |         k3, _ = model(X_nn + k2/2, hiddens)
61 |         k3 = h * k3.detach()
62 |         k4, _ = model(X_nn + k3, hiddens)
63 |         k4 = h * k4.detach()
64 |         X_nn = X_nn + 1/6*(k1 + 2*k2 + 2*k3 + k4)
65 | 
66 |         # Detach
67 |         X_nn = X_nn.detach()
68 |         hiddens = tuple(d.detach() for d in new_hiddens)
69 | 
70 |         logger.info(f"Timestep {i}")
71 | 
72 |         y = X_nn.cpu().numpy()
73 |         vel_error = np.sum((X_phy[i,:,n:] - y[0,:,n:])**2)
74 |         ang_error = (X_phy[i,:,:n] - y[0,:,:n])
75 |         while np.any(ang_error >= np.pi):
76 |             ang_error[ang_error >= np.pi] -= 2*np.pi
77 |         while np.any(ang_error < -np.pi):
78 |             ang_error[ang_error < -np.pi] += 2*np.pi
79 | 
80 |         ang_error = np.sum(ang_error**2)
81 |         errors[i] = (vel_error + ang_error)
82 | 
83 |     for i in range(args.steps):
84 |         print(f"{i}\t{np.sum(errors[0:i])}\t{errors[i]}")
85 | 
86 | 
87 | if __name__ == "__main__":
88 |     parser = argparse.ArgumentParser(description='Error of .')
89 |     parser.add_argument('--links', type=int, default=8, help="number of links")
90 |     parser.add_argument('--number', type=int, default=1000, help="number of starting positions to evaluate from")
91 |     parser.add_argument('--timestep', type=float, default=0.01, help="duration of each timestep")
92 | 
93 |     parser.add_argument('data', type=DynamicLoad("datasets"), help='the pendulum dataset to load the simulator from')
94 |     parser.add_argument('weight', type=latest_file, help='model weight to load')
95 |     parser.add_argument('steps', type=int, help="number of steps to evaluate over")
96 | 
97 |     main(parser.parse_args())
98 | 


--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | 
  3 | import argparse
  4 | import datetime
  5 | import glob
  6 | import os
  7 | import random
  8 | import sys
  9 | from collections import defaultdict
 10 | from pathlib import Path
 11 | 
 12 | import numpy as np
 13 | import torch
 14 | import torch.nn.functional as F
 15 | from tensorboardX import SummaryWriter
 16 | from torch import nn, optim
 17 | from torch.autograd import Variable
 18 | from torch.utils.data import DataLoader
 19 | from torchvision.utils import save_image
 20 | from util import DynamicLoad, setup_logging, to_variable
 21 | 
 22 | logger = setup_logging(os.path.basename(__file__))
 23 | 
 24 | def runbatch(args, model, loss, batch):
 25 |     X, Yactual = batch
 26 |     X = to_variable(X, cuda=torch.cuda.is_available())
 27 |     Yactual = to_variable(Yactual, cuda=torch.cuda.is_available())
 28 | 
 29 |     Ypred = model(X)
 30 |     return loss(Ypred, Yactual, X), Ypred
 31 | 
 32 | def test_model(args, model, test_dataloader, epoch=None, summarywriter=None):
 33 |     loss_parts = []
 34 |     model.eval()
 35 |     for batch_idx, data in enumerate(test_dataloader):
 36 |         loss, Ypred = runbatch(args, model, args.model.loss, data)
 37 |         loss_parts.append(np.array([l.cpu().item() for l in args.model.loss_flatten(loss)]))
 38 | 
 39 |         # Add parts to the summary if needed.
 40 |         args.model.summary(epoch, summarywriter, Ypred, data[0])
 41 | 
 42 |     return sum(loss_parts) / len(test_dataloader.dataset)
 43 | 
 44 | global _first_printed
 45 | _first_printed = False
 46 | def print_update(args, train_test, epoch, loss_elements):
 47 |     logger_progress = logger.getChild("progress")
 48 |     global _first_printed
 49 |     if not _first_printed:
 50 |         _first_printed = True
 51 |         loss_parts = "\t".join(args.model.loss_labels())
 52 |         s = f"TIMESTAMP\tTRAIN/TEST\tepoch\t{loss_parts}"
 53 |         logger_progress.info(s)
 54 |         print(s)
 55 | 
 56 |     now = datetime.datetime.now()
 57 |     loss_parts = "\t".join(map(str, loss_elements))
 58 |     s = f"{now}\t{train_test}\t{epoch}\t{loss_parts}"
 59 |     print(s)
 60 |     logger_progress.info(s)
 61 | 
 62 | def main(args):
 63 |     writer = SummaryWriter(logdir=args.log_to)
 64 |     model = args.model.model
 65 |     dataset = args.dataset
 66 | 
 67 |     train_dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
 68 |     if args.test_with:
 69 |         test_dataloader = DataLoader(args.test_with, batch_size=args.batch_size, shuffle=False)
 70 |     else:
 71 |         test_dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
 72 | 
 73 |     # TODO: Flag to disable CUDA
 74 |     if torch.cuda.is_available():
 75 |         model.cuda()
 76 | 
 77 |     optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
 78 | 
 79 |     # TODO: Resume training support
 80 |     for epoch in range(1, args.epochs + 1):
 81 |         model.train()
 82 |         loss_parts = []
 83 | 
 84 |         for batch_idx, data in enumerate(train_dataloader):
 85 |             optimizer.zero_grad()
 86 |             loss, _ = runbatch(args, model, args.model.loss, data)
 87 |             loss_parts.append(np.array([l.cpu().item() for l in args.model.loss_flatten(loss)]))
 88 | 
 89 |             optim_loss = loss[0] if isinstance(loss, (tuple, list)) else loss
 90 |             optim_loss.backward()
 91 |             optimizer.step()
 92 | 
 93 |         epoch_loss = sum(loss_parts) / len(dataset)
 94 |         print_update(args, "TRAIN", epoch, epoch_loss)
 95 |         for lbl, val in zip(args.model.loss_labels(), epoch_loss):
 96 |             writer.add_scalar(f'train_loss/{lbl}', val, epoch)
 97 | 
 98 |         if epoch % args.save_every == 0 or epoch == args.epochs:
 99 |             torch.save(model.state_dict(), args.weights.format(epoch=epoch))
100 |             test_loss = test_model(args, model, test_dataloader, epoch=epoch, summarywriter=writer)
101 |             print_update(args, "TEST", epoch, test_loss)
102 |             for lbl, val in zip(args.model.loss_labels(), test_loss):
103 |                 writer.add_scalar(f'test_loss/{lbl}', val, epoch)
104 | 
105 |     # Ensure the writer is completed.
106 |     writer.close()
107 | 
108 | if __name__ == "__main__":
109 |     parser = argparse.ArgumentParser(description='Train a VAE on a set of videos, fine-tune it on a single video, and generate the decoder.')
110 |     parser.set_defaults(func=lambda *a: parser.print_help())
111 | 
112 |     parser.add_argument('dataset', type=DynamicLoad("datasets"), help='dataset to train on')
113 |     parser.add_argument('model', type=DynamicLoad("models"), help='model to train with')
114 |     parser.add_argument('weights', type=str, help='save model weights')
115 |     parser.add_argument('--log-to', type=str, help='log destination within runs/')
116 |     parser.add_argument('--test-with', type=DynamicLoad("datasets"), default=None, help='dataset to test with instead of the training data')
117 |     parser.add_argument('--batch-size', type=int, default=256, help='batch size')
118 |     parser.add_argument('--learning-rate', type=float, default=1e-3, help='learning rate')
119 |     parser.add_argument('--epochs', type=int, default=120, help='number of epochs to run')
120 |     parser.add_argument('--save-every', type=int, default=100, help='save after this many epochs')
121 |     parser.set_defaults(func=main)
122 | 
123 |     try:
124 |         args = parser.parse_args()
125 |         main(args)
126 |     except:
127 |         raise
128 | 


--------------------------------------------------------------------------------
/train_lstm.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | 
  3 | import argparse
  4 | import datetime
  5 | import glob
  6 | import os
  7 | import random
  8 | import sys
  9 | from collections import defaultdict
 10 | from pathlib import Path
 11 | 
 12 | import numpy as np
 13 | import torch
 14 | import torch.nn.functional as F
 15 | from datasets import pendulum
 16 | from torch import nn, optim, tensor
 17 | from torch.autograd import Variable
 18 | from torch.utils.data import DataLoader
 19 | from torchvision.utils import save_image
 20 | from util import DynamicLoad, setup_logging, to_variable
 21 | 
 22 | logger = setup_logging(os.path.basename(__file__))
 23 | 
 24 | 
 25 | def trajectory_physics(number, n, steps):
 26 |     physics = pendulum.pendulum_gradient(n)
 27 | 
 28 |     h = 0.01 # Timestep
 29 | 
 30 |     cache_path = Path("pendulum-cache") / f"p-lstmtraj-{n}.npy"
 31 |     if not cache_path.exists():
 32 |         logger.info(f"Generating trajectories for {cache_path}")
 33 |         # Initialize args.number initial positions:
 34 |         X_init = np.zeros((number, 2*n)).astype(np.float32)
 35 |         X_init[:,:] = (np.random.rand(number, 2*n).astype(np.float32) - 0.5) * np.pi/2 # Pick values in range [-pi/4, pi/4] radians, radians/sec
 36 | 
 37 |         X_phy = np.zeros((steps, *X_init.shape), dtype=np.float32)
 38 |         X_grad = np.zeros((steps, *X_init.shape), dtype=np.float32)
 39 |         X_phy[0,...] = X_init
 40 |         X_grad[0,...] = physics(X_init)
 41 |         for i in range(1, steps):
 42 |             logger.info(f"Timestep {i}")
 43 |             k1 = h * physics(X_phy[i-1,...])
 44 |             k2 = h * physics(X_phy[i-1,...] + k1/2)
 45 |             k3 = h * physics(X_phy[i-1,...] + k2/2)
 46 |             k4 = h * physics(X_phy[i-1,...] + k3)
 47 |             X_phy[i,...] = X_phy[i-1,...] + 1/6*(k1 + 2*k2 + 2*k3 + k4)
 48 |             X_grad[i,...] = physics(X_phy[i,...])
 49 |             assert not np.any(np.isnan(X_phy[i,...]))
 50 | 
 51 |         np.save(cache_path, (X_phy, X_grad))
 52 |         logger.info(f"Done generating trajectories for {cache_path}")
 53 | 
 54 |     else:
 55 |         X_phy, X_grad = np.load(cache_path).astype(np.float32)
 56 |         logger.info(f"Loaded trajectories from {cache_path}. {X_phy.shape}, {X_grad.shape}")
 57 | 
 58 |     return X_phy, X_grad
 59 | 
 60 | N = 8
 61 | STEPS = 50
 62 | TRAJECTORIES = 32
 63 | 
 64 | class LSTMModel(nn.Module):
 65 |     def __init__(self, embedding_dim, hidden_dim=116):
 66 |         super().__init__()
 67 |         self.hidden_dim = hidden_dim
 68 | 
 69 |         # The LSTM takes word embeddings as inputs, and outputs hidden states
 70 |         # with dimensionality hidden_dim.
 71 |         self.enc = nn.Sequential(nn.Linear(embedding_dim, hidden_dim), nn.ReLU())
 72 |         self.lstm = nn.LSTM(hidden_dim, hidden_dim)
 73 |         self.dec = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, embedding_dim))
 74 | 
 75 |     def forward(self, x, hidden):
 76 |         x = self.enc(x)
 77 |         x, hidden = self.lstm(x, hidden)
 78 |         return self.dec(x), hidden
 79 | 
 80 | 
 81 | def main(args):
 82 |     hidden_dim = 116
 83 |     N = args.links
 84 |     model = LSTMModel(2*N, hidden_dim)
 85 | 
 86 |     X_phy, X_grad = trajectory_physics(TRAJECTORIES, N, STEPS)
 87 | 
 88 |     X_phy = tensor(X_phy)
 89 |     X_grad = tensor(X_grad)
 90 | 
 91 |     # TODO: Flag to disable CUDA
 92 |     if torch.cuda.is_available():
 93 |         model.cuda()
 94 |         X_phy = X_phy.cuda()
 95 |         X_grad = X_grad.cuda()
 96 | 
 97 |     optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
 98 | 
 99 |     for epoch in range(1, args.epochs + 1):
100 |         model.train()
101 |         loss_parts = []
102 | 
103 |         hiddens = None
104 |         for timestep in range(X_phy.shape[0]):
105 |             optimizer.zero_grad()
106 | 
107 |             y, new_hiddens = model(X_phy[timestep,...].unsqueeze(0), hiddens)
108 |             hiddens = tuple(d.detach() for d in new_hiddens)
109 |             loss = torch.sum((y - X_grad[timestep,...])**2)
110 |             loss_parts.append(loss.cpu().item())
111 | 
112 |             loss.backward()
113 |             optimizer.step()
114 | 
115 |         epoch_loss = sum(loss_parts) / X_phy.shape[0]
116 |         print(f"TRAIN\t{epoch}\t{epoch_loss}")
117 | 
118 |         if epoch % args.save_every == 0 or epoch == args.epochs:
119 |             torch.save(model.state_dict(), args.weights.format(epoch=epoch))
120 | 
121 | 
122 | if __name__ == "__main__":
123 |     parser = argparse.ArgumentParser(description='Train a VAE on a set of videos, fine-tune it on a single video, and generate the decoder.')
124 |     parser.set_defaults(func=lambda *a: parser.print_help())
125 | 
126 |     parser.add_argument('weights', type=str, help='save model weights')
127 | 
128 |     parser.add_argument('--links', type=int, default=8, help="number of links")
129 |     parser.add_argument('--log-to', type=str, help='log destination within runs/')
130 |     parser.add_argument('--learning-rate', type=float, default=5 e-4, help='learning rate')
131 |     parser.add_argument('--epochs', type=int, default=500, help='number of epochs to run')
132 |     parser.add_argument('--save-every', type=int, default=100, help='save after this many epochs')
133 |     parser.set_defaults(func=main)
134 | 
135 |     try:
136 |         args = parser.parse_args()
137 |         main(args)
138 |     except:
139 |         raise
140 | 


--------------------------------------------------------------------------------
/train_many:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | 
 3 | for ALPHA in 0.001 0.0001
 4 | do
 5 | for INNER_EPSILON in 0.01 0.005 0.001
 6 | do
 7 | for SMOOTH_V in 0
 8 | do
 9 | for REHU in 0.005 0.001 0.0005
10 | do
11 | ./train_pendulum_stable stable "$1" "$ALPHA" NN-REHU "$INNER_EPSILON" "$SMOOTH_V" "$REHU" &
12 | done
13 | wait
14 | done
15 | done
16 | done
17 | 


--------------------------------------------------------------------------------
/train_pendulum_simple:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | set -e
 3 | 
 4 | N="$1"
 5 | 
 6 | # Latent space dim:
 7 | LSD=$(( 2 * N ))
 8 | 
 9 | OUTDIR="experiments/pendulum-simple/${N}"
10 | MODEL="simple[a=$LSD,b=116]"
11 | 
12 | mkdir -p "$OUTDIR"
13 | echo $MODEL > "$OUTDIR/model"
14 | 
15 | date >> "$OUTDIR/progress.txt"
16 | ./.colorize ./train.py                      \
17 |     --log-to "runs/$OUTDIR"                 \
18 |     --batch-size 2000                       \
19 |     --learning-rate "0.001"                 \
20 |     --epochs 1000                           \
21 |     --test-with "pendulum[n=$N,test]"       \
22 |     "pendulum[n=$N]"                        \
23 |     "$MODEL"                                \
24 |     "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt"
25 | 
26 | ./.colorize ./pendulum_error.py             \
27 |     pendulum[n=$N,test]                     \
28 |     "$MODEL"                                \
29 |     "$OUTDIR/checkpoint_*.pth"    \
30 |     1000 | tee -a "$OUTDIR/eval.txt"
31 | 


--------------------------------------------------------------------------------
/train_pendulum_stable:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | set -e
 3 | 
 4 | EXP_ID="$1"
 5 | N="$2"
 6 | ALPHA="$3"
 7 | INNER="$4"
 8 | INNER_EPSILON="$5"
 9 | SMOOTH_V="$6"
10 | REHU="$7"
11 | 
12 | # Latent space dim:
13 | LSD=$(( 2 * N ))
14 | 
15 | OUTDIR="experiments/pendulum-${EXP_ID}/${N}_${ALPHA}_${INNER}_${INNER_EPSILON}_${SMOOTH_V}_${REHU}"
16 | MODEL="stabledynamics[latent_space_dim=$LSD,a=$ALPHA,projfn=$INNER,projfn_eps=$INNER_EPSILON,smooth_v=$SMOOTH_V,hp=60,h=100,rehu=$REHU]"
17 | 
18 | mkdir -p "$OUTDIR"
19 | echo $MODEL > "$OUTDIR/model"
20 | 
21 | date >> "$OUTDIR/progress.txt"
22 | ./.colorize ./train.py                      \
23 |     --log-to "runs/$OUTDIR"                 \
24 |     --batch-size 2000                       \
25 |     --learning-rate "0.001"                 \
26 |     --epochs 1000                           \
27 |     --test-with "pendulum[n=$N,test]"       \
28 |     "pendulum[n=$N]"                        \
29 |     "$MODEL"                                \
30 |     "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt"
31 | 
32 | ./.colorize ./pendulum_error.py             \
33 |     pendulum[n=$N,test]                     \
34 |     "$MODEL"                                \
35 |     "$OUTDIR/checkpoint_*.pth"              \
36 |     1000 | tee -a "$OUTDIR/eval.txt"
37 | 


--------------------------------------------------------------------------------
/train_pendulum_stable_knownV:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | set -e
 3 | 
 4 | EXP_ID="$1"
 5 | N="$2"
 6 | ALPHA="$3"
 7 | INNER="ActualPendulumEnergy"
 8 | INNER_EPSILON="$5"
 9 | SMOOTH_V="$6"
10 | REHU="$7"
11 | 
12 | # Latent space dim:
13 | LSD=$(( 2 * N ))
14 | 
15 | OUTDIR="experiments/pendulum-${EXP_ID}/${N}_${ALPHA}_${INNER}_${INNER_EPSILON}_${SMOOTH_V}_${REHU}"
16 | MODEL="stabledynamics[latent_space_dim=$LSD,a=$ALPHA,projfn=$INNER,pendulum_n=$N,smooth_v=$SMOOTH_V,hp=60,h=100,rehu=$REHU]"
17 | 
18 | mkdir -p "$OUTDIR"
19 | echo $MODEL > "$OUTDIR/model"
20 | 
21 | date >> "$OUTDIR/progress.txt"
22 | ./.colorize ./train.py                      \
23 |     --log-to "runs/$OUTDIR"                 \
24 |     --batch-size 2000                       \
25 |     --learning-rate "0.005"                 \
26 |     --epochs 1000                           \
27 |     --test-with "pendulum[n=$N,test]"       \
28 |     "pendulum[n=$N]"                        \
29 |     "$MODEL"                                \
30 |     "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt"
31 | 
32 | ./.colorize ./pendulum_error.py             \
33 |     pendulum[n=$N,test]                     \
34 |     "$MODEL"                                \
35 |     "$OUTDIR/checkpoint_*.pth"              \
36 |     1000 | tee -a "$OUTDIR/eval.txt"
37 | 


--------------------------------------------------------------------------------
/train_vae_simple:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | 
 3 | DATASET="$1"
 4 | LR="0.0005"
 5 | W1="-0.25"
 6 | INNER="PSICNN"
 7 | 
 8 | if [ -z "$DATASET" ]; then
 9 |     echo "NO DATASET"
10 |     exit
11 | fi
12 | 
13 | OUTDIR="experiments/vae_simple/${DATASET}"
14 | mkdir -p "$OUTDIR"
15 | MODEL="vae_trajectory[w=$W1,projfn=$INNER]"
16 | 
17 | date >> "$OUTDIR/progress.txt"
18 | ./.colorize ./train.py                      \
19 |     --log-to "runs/$OUTDIR"                 \
20 |     --batch-size 300                        \
21 |     --learning-rate "$LR"                   \
22 |     --epochs 150                            \
23 |     imagepairs[files=youtube/$DATASET/*]    \
24 |     "$MODEL"                                \
25 |     "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt"
26 | 
27 | ./render_vae "$DATASET" "$MODEL" "$OUTDIR"
28 | 


--------------------------------------------------------------------------------
/train_vae_stable:
--------------------------------------------------------------------------------
 1 | #!/bin/bash
 2 | 
 3 | set -e
 4 | 
 5 | EXPERIMENT_ID="$1"
 6 | DATASET="$2"
 7 | W1="$3"
 8 | W2="$4"
 9 | INNER="$5"
10 | ALPHA="0.0001"
11 | REHU="0.0001"
12 | SCALE_FX="$6"
13 | 
14 | if [ -z "$EXPERIMENT_ID" ]; then
15 |     echo "NO EXPERIMENT ID"
16 |     exit
17 | elif [ -z "$DATASET" ]; then
18 |     echo "NO DATASET"
19 |     exit
20 | fi
21 | 
22 | OUTDIR="experiments/${EXPERIMENT_ID}/${DATASET}_${W1}_${W2}_${INNER}_${ALPHA}_${REHU}_scaled${SCALE_FX}"
23 | MODEL="vae_trajectory[stable,w=$W1,bce=$W2,a=$ALPHA,projfn=$INNER,h=320,hp=200,rehu=$REHU,scale_fx=${SCALE_FX}]"
24 | 
25 | mkdir -p "$OUTDIR"
26 | echo $MODEL > "$OUTDIR/model"
27 | 
28 | date >> "$OUTDIR/progress.txt"
29 | ./.colorize ./train.py                      \
30 |     --log-to "runs/$OUTDIR"                 \
31 |     --batch-size 260                        \
32 |     --learning-rate "0.0002"                \
33 |     --epochs 150                            \
34 |     imagepairs[files=youtube/$DATASET/*]    \
35 |     "$MODEL"                                \
36 |     "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt"
37 | 
38 | ./render_vae "$DATASET" "$MODEL" "$OUTDIR"
39 | 


--------------------------------------------------------------------------------
/trajectory.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import argparse
 4 | import datetime
 5 | import glob
 6 | import os
 7 | import random
 8 | import sys
 9 | from collections import defaultdict
10 | from pathlib import Path
11 | 
12 | import h5py
13 | import numpy as np
14 | import torch
15 | import torch.nn.functional as F
16 | from torch import nn, optim
17 | from torch.autograd import Variable
18 | from torch.utils.data import DataLoader
19 | from torchvision.utils import save_image
20 | from util import DynamicLoad, latest_file, setup_logging, to_variable
21 | 
22 | logger = setup_logging(os.path.basename(__file__))
23 | 
24 | def main(args):
25 |     model = args.model.model
26 |     model.load_state_dict(torch.load(args.weight))
27 | 
28 |     dataset = args.dataset
29 |     test_dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
30 | 
31 |     # TODO: Flag to disable CUDA
32 |     if torch.cuda.is_available():
33 |         model.cuda()
34 | 
35 |     model.eval()
36 |     trajectory = []
37 |     for batch_idx, data in enumerate(test_dataloader):
38 |         img, _ = data
39 |         img = to_variable(img, cuda=torch.cuda.torch.cuda.is_available())
40 | 
41 |         # z = model(img)
42 |         # Break the abstraction here:
43 |         # Y_a, mu_a, logvar_a, z_a, Y_b, z_b = model(img)
44 |         (Y_a, Code_a, Quantized_a, Perplexity_a, Latent_a, Y_b, Quantized_b, Latent_b) = model(img)
45 |         if args.mu:
46 |             z = mu_a
47 |         else:
48 |             # z = z_a
49 |             z = Latent_a
50 | 
51 |         trajectory.append(z.cpu().data.numpy())
52 | 
53 |     trajectory = np.concatenate(trajectory)
54 |     h5f = h5py.File(args.output, 'w')
55 |     h5f.create_dataset('seq', data=[trajectory])
56 |     h5f.create_dataset('param', data=[])
57 |     h5f.close()
58 | 
59 |     mn = np.mean(trajectory, axis=0)
60 |     linf = np.max(np.abs(trajectory - mn))
61 | 
62 |     print(f"Data: {trajectory.shape}; Mean: {mn.shape}; Linf: {linf}")
63 | 
64 |     np.savetxt(Path(args.output).with_suffix(".mean"), mn)
65 |     Path(args.output).with_suffix(".linf").write_text(str(linf))
66 | 
67 | 
68 | if __name__ == "__main__":
69 |     parser = argparse.ArgumentParser(description='Train a VAE on a set of videos, fine-tune it on a single video, and generate the decoder.')
70 |     parser.set_defaults(func=lambda *a: parser.print_help())
71 | 
72 |     parser.add_argument('dataset', type=DynamicLoad("datasets"), help='dataset to train on')
73 |     parser.add_argument('model', type=DynamicLoad("models"), help='model to train with')
74 |     parser.add_argument('weight', type=latest_file, help='save model weights')
75 |     parser.add_argument('output', type=str, help='trajectory file')
76 |     parser.add_argument('--test-with', type=DynamicLoad("datasets"), help='dataset to test with instead of the training data')
77 |     parser.add_argument('--batch-size', type=int, default=256, help='batch size')
78 |     parser.add_argument('--mu', action="store_true", help='sample with zero variance')
79 |     parser.set_defaults(func=main)
80 | 
81 |     try:
82 |         args = parser.parse_args()
83 |         main(args)
84 |     except:
85 |         raise
86 | 


--------------------------------------------------------------------------------
/util.py:
--------------------------------------------------------------------------------
  1 | import glob
  2 | import importlib
  3 | import logging
  4 | import re
  5 | 
  6 | import h5py
  7 | from torch.autograd import Variable
  8 | 
  9 | PROJECT="videotexture"
 10 | 
 11 | def to_variable(X, cuda=True):
 12 |     if isinstance(X, (tuple, list)):
 13 |         return tuple(to_variable(x) for x in X)
 14 |     else:
 15 |         X = Variable(X)
 16 |         if cuda:
 17 |             return X.cuda().requires_grad_()
 18 |         return X.requires_grad_()
 19 | 
 20 | def setup_logging(name=None):
 21 |     if name:
 22 |         name = f"{PROJECT}.{name}"
 23 |     else:
 24 |         name = PROJECT
 25 | 
 26 |     logfile = logging.FileHandler(filename=f"{name}.log", mode='w')
 27 |     logfile.setLevel(logging.DEBUG)
 28 |     logfile.setFormatter(logging.Formatter('[%(asctime)s, %(levelname)s @%(name)s] %(message)s'))
 29 |     console = logging.StreamHandler()
 30 |     console.setLevel(logging.INFO)
 31 |     console.setFormatter(logging.Formatter('[%(asctime)s %(levelname)-3s @%(name)s] %(message)s', datefmt='%H:%M:%S'))
 32 |     logging.basicConfig(level=logging.DEBUG, handlers=[logfile, console])
 33 |     return logging.getLogger(name)
 34 | 
 35 | def latest_file(globstr):
 36 |     fn_candidates = sorted(glob.glob(globstr), reverse=True)
 37 |     if len(fn_candidates) == 0:
 38 |         raise RuntimeError(f"No files match the pattern {globstr}")
 39 |     return fn_candidates[0]
 40 | 
 41 | def DynamicLoad(pkg):
 42 |     logger = setup_logging("DynamicLoad")
 43 |     # Used to dynamically load modules in commandline options.
 44 | 
 45 |     def _split_name(modnamefull):
 46 |         m = re.fullmatch('(\w+)(\[((\w+(=[\w\.\+\-~@/_\*]+)?,?)*)\])?', modnamefull)
 47 |         if m is None:
 48 |             logger.error(f"Incorrect import string {modnamefull}")
 49 |             raise RuntimeError("Incorrect import.")
 50 | 
 51 |         modname = m.group(1)
 52 |         props = m.group(3)
 53 |         if props is not None:
 54 |             # Convert it to a dictionary:
 55 |             props = (v.split('=', 1) for v in props.split(","))
 56 |             props = {v[0]: (v[1] if len(v) == 2 else True) for v in props}
 57 |         else:
 58 |             props = {}
 59 | 
 60 |         return modname, props
 61 | 
 62 |     def _load_actual(modnamefull):
 63 |         modname, props = _split_name(modnamefull)
 64 |         try:
 65 |             load_mod = importlib.import_module(pkg + "." + modname, package="")
 66 |         except Exception as e:
 67 |             logger.exception(e)
 68 |             raise
 69 | 
 70 |         try:
 71 |             logger.debug("Configuring {}.{} with options: {}".format(pkg, modname, str(props)))
 72 |             if hasattr(load_mod, 'build'):
 73 |                 return load_mod.build(props)
 74 |             elif hasattr(load_mod, 'configure'):
 75 |                 load_mod.configure(props)
 76 |                 return load_mod
 77 |             elif len(props) == 0:
 78 |                 return load_mod
 79 |             else:
 80 |                 raise RuntimeError('Options passed to item that does not support options.')
 81 | 
 82 |         except Exception as e:
 83 |             logger.exception(e)
 84 |             raise e
 85 | 
 86 |     return _load_actual
 87 | 
 88 | def loadDataFile(fn):
 89 |     h5f = h5py.File(fn)
 90 |     seq = []
 91 | 
 92 |     in_seq = h5f['seq']
 93 |     if isinstance(in_seq, h5py.Dataset):
 94 |         seq = in_seq
 95 |     else:
 96 |         for v in in_seq.values():
 97 |             seq.append(v[()])
 98 | 
 99 |     params = None
100 |     if 'param' in h5f:
101 |         params = h5f['param'][()]
102 |     return (seq, params)
103 | 


--------------------------------------------------------------------------------
/vae_trajmodel.py:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python3
 2 | 
 3 | import argparse
 4 | import os
 5 | 
 6 | import matplotlib.cm as cm
 7 | import matplotlib.pyplot as plt
 8 | import numpy as np
 9 | import torch
10 | from plot_data import plot_data, plot_data_args
11 | from torchvision.utils import save_image
12 | from util import (DynamicLoad, latest_file, loadDataFile, setup_logging,
13 |                   to_variable)
14 | 
15 | logger = setup_logging(os.path.basename(__file__))
16 | 
17 | def main(args):
18 |     model = args.model.model
19 |     model.load_state_dict(torch.load(args.weight))
20 |     model.eval()
21 |     if torch.cuda.is_available():
22 |         model.cuda()
23 | 
24 |     dyn_model = model.dyn
25 |     vae_model = model.vae
26 | 
27 |     # Load sequence:
28 |     seq, param = args.data
29 |     original_seq = seq[args.select,:,:]
30 |     trajectory = [original_seq[args.start_step:(args.start_step+1),:]]
31 |     renders = []
32 | 
33 |     get_image = lambda r: torch.squeeze(vae_model.decode(r).detach().cpu(), dim=0)
34 | 
35 |     # Start with the first index:
36 |     X = to_variable(torch.tensor(trajectory[0]), cuda=torch.cuda.is_available())
37 |     X.requires_grad_()
38 |     if args.act == "render":
39 |         renders.append(get_image(X))
40 | 
41 |     for i in range(args.steps):
42 |         X = to_variable((X + dyn_model(X)).data, cuda=torch.cuda.is_available())
43 |         X.requires_grad_()
44 | 
45 |         trajectory.append(X.detach().cpu().numpy())
46 |         if args.act == "render":
47 |             renders.append(get_image(X))
48 | 
49 |     trajectory = np.squeeze(np.stack(trajectory), axis=1)
50 | 
51 |     if args.act == "plot":
52 |         plot_data(args, trajectory, original_seq=original_seq)
53 |     elif args.act == "render":
54 |         if args.save:
55 |             save_image(renders, args.save)
56 |         if args.save_frames:
57 |             for i, im in enumerate(renders):
58 |                 save_image(im, args.save_frames.format(i))
59 | 
60 | 
61 | if __name__ == "__main__":
62 |     parser = argparse.ArgumentParser(description='Plot data files.')
63 |     parser.add_argument('--select', type=int, default=0, help="select the starting position from this index")
64 |     parser.add_argument('--start-step', type=int, default=0, help="number of steps to render")
65 |     parser.add_argument('--steps', type=int, default=240, help="number of steps to render")
66 | 
67 |     parser.add_argument('data', type=loadDataFile, help='the data file to load the start position from')
68 |     parser.add_argument('model', type=DynamicLoad("models"), help='model to load')
69 |     parser.add_argument('weight', type=latest_file, help='model weight to load')
70 | 
71 |     subparsers = parser.add_subparsers(help='output to produce')
72 | 
73 |     parser_plot = plot_data_args(subparsers.add_parser('plot', help='plot trajectories in latent space'))
74 |     parser_plot.set_defaults(act="plot")
75 | 
76 |     parser_render = subparsers.add_parser('render', help='create images from the latent space')
77 |     parser_render.set_defaults(act="render")
78 |     parser_render.add_argument('--save', type=str, help="save renders as a single image")
79 |     parser_render.add_argument('--save-frames', type=str, help="save frames, pass a format string")
80 | 
81 |     parser_field = subparsers.add_parser('field', help='create a field plot of the latent space')
82 |     parser_field.set_defaults(act="field")
83 |     parser_field.add_argument('--save', type=str, help="save field plot as an image")
84 | 
85 |     main(parser.parse_args())
86 | 


--------------------------------------------------------------------------------