93 | |
94 |
95 | ```py
96 | class PianoGenieDecoder(nn.Module):
97 | def __init__(self, rnn_dim=128, rnn_num_layers=2):
98 | super().__init__()
99 | self.rnn_dim = rnn_dim
100 | self.rnn_num_layers = rnn_num_layers
101 | self.input = nn.Linear(PIANO_NUM_KEYS + 3, rnn_dim)
102 | self.lstm = nn.LSTM(
103 | rnn_dim,
104 | rnn_dim,
105 | rnn_num_layers,
106 | batch_first=True,
107 | bidirectional=False,
108 | )
109 | self.output = nn.Linear(rnn_dim, 88)
110 |
111 | def init_hidden(self, batch_size, device=None):
112 | h = torch.zeros(self.rnn_num_layers, batch_size, self.rnn_dim, device=device)
113 | c = torch.zeros(self.rnn_num_layers, batch_size, self.rnn_dim, device=device)
114 | return (h, c)
115 |
116 | def forward(self, k_m1, t, b, h_0=None):
117 | # Encode input
118 | inputs = [
119 | F.one_hot(k_m1, PIANO_NUM_KEYS + 1),
120 | t.unsqueeze(dim=2),
121 | b.unsqueeze(dim=2),
122 | ]
123 | x = torch.cat(inputs, dim=2)
124 |
125 | # Project encoded inputs
126 | x = self.input(x)
127 |
128 | # Run RNN
129 | if h_0 is None:
130 | h_0 = self.init_hidden(k.shape[0], device=k.device)
131 | x, h_N = self.lstm(x, h_0)
132 |
133 | # Compute logits
134 | hat_k = self.output(x)
135 |
136 | return hat_k, h_N
137 | ```
138 |
139 | |
140 |
141 |
142 | ```js
143 | class PianoGenieDecoder extends Module {
144 | constructor(rnnDim, rnnNumLayers) {
145 | super();
146 | this.rnnDim = rnnDim === undefined ? 128 : rnnDim;
147 | this.rnnNumLayers = rnnNumLayers === undefined ? 2 : rnnNumLayers;
148 | }
149 |
150 | initHidden(batchSize) {
151 | // NOTE: This allocates memory that must later be freed
152 | const c = [];
153 | const h = [];
154 | for (let i = 0; i < this.rnnNumLayers; ++i) {
155 | c.push(tf.zeros([batchSize, this.rnnDim], "float32"));
156 | h.push(tf.zeros([batchSize, this.rnnDim], "float32"));
157 | }
158 | return new LSTMHiddenState(c, h);
159 | }
160 |
161 | forward(kim1, ti, bi, him1) {
162 | // Encode input
163 | const inputs = [
164 | tf.oneHot(kim1, PIANO_NUM_KEYS + 1),
165 | tf.expandDims(ti, 1),
166 | tf.expandDims(bi, 1)
167 | ];
168 | let x = tf.concat(inputs, 1);
169 |
170 | // Project encoded inputs
171 | x = tf.add(
172 | tf.matMul(x, this._params["dec.input.weight"], false, true),
173 | this._params[`dec.input.bias`]
174 | );
175 |
176 | // Run RNN
177 | if (him1 === undefined || him1 === null) {
178 | him1 = this.initHidden(kim1.shape[0]);
179 | }
180 | const [hic, hih] = tf.multiRNNCell(this._cells, x, him1.c, him1.h);
181 | x = hih[this.rnnNumLayers - 1];
182 | const hi = new LSTMHiddenState(hic, hih);
183 |
184 | // Compute logits
185 | const hatki = tf.add(
186 | tf.matMul(x, this._params["dec.output.weight"], false, true),
187 | this._params[`dec.output.bias`]
188 | );
189 |
190 | return [hatki, hi];
191 | }
192 | }
193 | ```
194 |
195 | |
196 |
197 | 
![](\"https://raw.githubusercontent.com/chrisdonahue/music-cocreation-tutorial/main/part-1-py-training/figures/overview.png\")
![](\"https://raw.githubusercontent.com/chrisdonahue/music-cocreation-tutorial/main/part-1-py-training/figures/decoder.png\")
![](\"https://i.imgur.com/P3bQFsC.png\")
![](\"https://raw.githubusercontent.com/chrisdonahue/music-cocreation-tutorial/main/part-1-py-training/figures/quantization.png\")