forked from salu133445/musegan
-
Notifications
You must be signed in to change notification settings - Fork 1
/
config.py
338 lines (290 loc) · 13.4 KB
/
config.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
"""Define configuration variables in experiment, model and training levels.
Quick Setup
===========
Change the values in the dictionary `SETUP` for a quick setup.
Documentation is provided right after each key.
Configuration
=============
More configuration options are providedin as a dictionary `CONFIG`.
`CONFIG['exp']`, `CONFIG['data']`, `CONFIG['model']`, `CONFIG['train']` and
`CONFIG['tensorflow']` define experiment-, data-, model-, training-,
TensorFlow-related configuration variables, respectively.
Note that the automatically-determined experiment name is based only on the
values defined in the dictionary `SETUP`, so remember to provide the experiment
name manually if you have changed the configuration so that you won't overwrite
existing experiment directories.
"""
import os
import shutil
import distutils.dir_util
import importlib
import numpy as np
import tensorflow as tf
# Quick setup
SETUP = {
'model': 'musegan',
# {'musegan', 'bmusegan'}
# The model to use. Currently support MuseGAN and BinaryMuseGAN models.
'exp_name': None,
# The experiment name. A folder with the same name will be created in 'exp/'
# directory where all the experiment-related files will be saved. None to
# determine automatically. Note that the automatically-determined experiment
# name is based only on the values defined in the dictionary `SETUP`, so
# remember to provide the experiment name manually when you modify any other
# configuration variable (so that you won't overwrite a trained model).
'prefix': 'lastfm_alternative',
# Prefix for the experiment name. Useful when training with different
# training data to avoid replacing the previous experiment outputs.
'training_data': 'lastfm_alternative_8b_phrase.npy',
# Path to the training data. The training data can be loaded from a npy
# file in the hard disk or from the shared memory using SharedArray package.
# Note that the data will be reshaped to (-1, num_bar, num_timestep,
# num_pitch, num_track) and remember to set these variable to proper values,
# which are defined in `CONFIG['model']`.
'training_data_location': 'hd',
# Location of the training data. 'hd' to load from a npy file stored in the
# hard disk. 'sa' to load from shared memory using SharedArray package.
'gpu': '0',
# The GPU index in os.environ['CUDA_VISIBLE_DEVICES'] to use.
'preset_g': 'hybrid',
# MuseGAN: {'composer', 'jamming', 'hybrid'}
# BinaryMuseGAN: {'proposed', 'proposed_small'}
# Use a preset network architecture for the generator or set to None and
# setup `CONFIG['model']['net_g']` to define the network architecture.
'preset_d': 'proposed',
# {'proposed', 'proposed_small', 'ablated', 'baseline', None}
# Use a preset network architecture for the discriminator or set to None
# and setup `CONFIG['model']['net_d']` to define the network architecture.
'pretrained_dir': None,
# The directory containing the pretrained model. None to retrain the
# model from scratch.
'verbose': True,
# True to print each batch details to stdout. False to print once an epoch.
'sample_along_training': True,
# True to generate samples along the training process. False for nothing.
'evaluate_along_training': True,
# True to run evaluation along the training process. False for nothing.
# ------------------------- For BinaryMuseGAN only -------------------------
'two_stage_training': True,
# True to train the model in a two-stage training setting. False to
# train the model in an end-to-end manner.
'training_phase': 'first_stage',
# {'first_stage', 'second_stage'}
# The training phase in a two-stage training setting. Only effective
# when `two_stage_training` is True.
'first_stage_dir': None,
# The directory containing the pretrained first-stage model. None to
# determine automatically (assuming using default `exp_name`). Only
# effective when two_stage_training` is True and `training_phase` is
# 'second_stage'.
'joint_training': False,
# True to train the generator and the refiner jointly. Only effective
# when `two_stage_training` is True and `training_phase` is 'second_stage'.
'preset_r': 'proposed_bernoulli',
# {'proposed_round', 'proposed_bernoulli'}
# Use a preset network architecture for the refiner or set to None and
# setup `CONFIG['model']['net_r']` to define the network architecture.
}
CONFIG = {}
#===============================================================================
#=========================== TensorFlow Configuration ==========================
#===============================================================================
os.environ['CUDA_VISIBLE_DEVICES'] = SETUP['gpu']
CONFIG['tensorflow'] = tf.ConfigProto()
CONFIG['tensorflow'].gpu_options.allow_growth = True
#===============================================================================
#========================== Experiment Configuration ===========================
#===============================================================================
CONFIG['exp'] = {
'model': None,
'exp_name': None,
'pretrained_dir': None,
'two_stage_training': None, # For BinaryMuseGAN only
'first_stage_dir': None, # For BinaryMuseGAN only
}
for key in ('model', 'pretrained_dir'):
if CONFIG['exp'][key] is None:
CONFIG['exp'][key] = SETUP[key]
if SETUP['model'] == 'musegan':
# Set default experiment name
if CONFIG['exp']['exp_name'] is None:
if SETUP['exp_name'] is not None:
CONFIG['exp']['exp_name'] = SETUP['exp_name']
else:
CONFIG['exp']['exp_name'] = '_'.join(
(SETUP['prefix'], 'g', SETUP['preset_g'], 'd',
SETUP['preset_d']))
if SETUP['model'] == 'bmusegan':
if CONFIG['exp']['two_stage_training'] is None:
CONFIG['exp']['two_stage_training'] = SETUP['two_stage_training']
# Set default experiment name
if CONFIG['exp']['exp_name'] is None:
if SETUP['exp_name'] is not None:
CONFIG['exp']['exp_name'] = SETUP['exp_name']
elif not SETUP['two_stage_training']:
CONFIG['exp']['exp_name'] = '_'.join(
(SETUP['prefix'], 'end2end', 'g', SETUP['preset_g'], 'd',
SETUP['preset_d'], 'r', SETUP['preset_r']))
elif SETUP['training_phase'] == 'first_stage':
CONFIG['exp']['exp_name'] = '_'.join(
(SETUP['prefix'], SETUP['training_phase'], 'g',
SETUP['preset_g'], 'd', SETUP['preset_d']))
elif SETUP['training_phase'] == 'second_stage':
if SETUP['joint_training']:
CONFIG['exp']['exp_name'] = '_'.join(
(SETUP['prefix'], SETUP['training_phase'], 'joint', 'g',
SETUP['preset_g'], 'd', SETUP['preset_d'], 'r',
SETUP['preset_r']))
else:
CONFIG['exp']['exp_name'] = '_'.join(
(SETUP['prefix'], SETUP['training_phase'], 'g',
SETUP['preset_g'], 'd', SETUP['preset_d'], 'r',
SETUP['preset_r']))
# Set default first stage model directory
if CONFIG['exp']['first_stage_dir'] is None:
if SETUP['first_stage_dir'] is not None:
CONFIG['exp']['first_stage_dir'] = SETUP['first_stage_dir']
else:
CONFIG['exp']['first_stage_dir'] = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'exp',
'_'.join((SETUP['prefix'], 'first_stage', 'g',
SETUP['preset_g'], 'd', SETUP['preset_d'])),
'checkpoints')
#===============================================================================
#============================= Data Configuration ==============================
#===============================================================================
CONFIG['data'] = {
'training_data': None,
'training_data_location': None,
}
for key in ('training_data', 'training_data_location'):
if CONFIG['data'][key] is None:
CONFIG['data'][key] = SETUP[key]
#===============================================================================
#=========================== Training Configuration ============================
#===============================================================================
CONFIG['train'] = {
'num_epoch': 10,
'verbose': None,
'sample_along_training': None,
'evaluate_along_training': None,
'two_stage_training': None, # For BinaryMuseGAN only
'training_phase': None, # For BinaryMuseGAN only
'slope_annealing_rate': 1.1, # For BinaryMuseGAN only
}
for key in ('verbose', 'sample_along_training', 'evaluate_along_training'):
if CONFIG['train'][key] is None:
CONFIG['train'][key] = SETUP[key]
if SETUP['model'] == 'bmusegan' and CONFIG['train']['training_phase'] is None:
CONFIG['train']['training_phase'] = SETUP['training_phase']
#===============================================================================
#============================= Model Configuration =============================
#===============================================================================
CONFIG['model'] = {
'joint_training': None, # For BinaryMuseGAN only
# Parameters
'batch_size': 32, # Note: tf.layers.conv3d_transpose requires a fixed batch
# size in TensorFlow < 1.6
'gan': {
'type': 'wgan-gp', # 'gan', 'wgan', 'wgan-gp'
'clip_value': .01,
'gp_coefficient': 10.
},
'optimizer': {
# Parameters for the Adam optimizers
'lr': .002,
'beta1': .5,
'beta2': .9,
'epsilon': 1e-8
},
# Data
'num_bar': 4,
'num_beat': 4,
'num_pitch': 84,
'num_track': 8,
'num_timestep': 96,
'beat_resolution': 24,
'lowest_pitch': 24, # MIDI note number of the lowest pitch in data tensors
# Tracks
'track_names': (
'Drums', 'Piano', 'Guitar', 'Bass', 'Ensemble', 'Reed', 'Synth Lead',
'Synth Pad'
),
'programs': (0, 0, 24, 32, 48, 64, 80, 88),
'is_drums': (True, False, False, False, False, False, False, False),
# Network architectures (define them here if not using the presets)
'net_g': None,
'net_d': None,
'net_r': None, # For BinaryMuseGAN only
# Playback
'pause_between_samples': 96,
'tempo': 90.,
# Samples
'num_sample': 16,
'sample_grid': (2, 8),
# Metrics
'metric_map': np.array([
# indices of tracks for the metrics to compute
[True] * 8, # empty bar rate
[True] * 8, # number of pitch used
[False] + [True] * 7, # qualified note rate
[False] + [True] * 7, # polyphonicity
[False] + [True] * 7, # in scale rate
[True] + [False] * 7, # in drum pattern rate
[False] + [True] * 7 # number of chroma used
], dtype=bool),
'tonal_distance_pairs': [(1, 2)], # pairs to compute the tonal distance
'scale_mask': list(map(bool, [1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1])),
'drum_filter': np.tile([1., .1, 0., 0., 0., .1], 16),
'tonal_matrix_coefficient': (1., 1., .5),
# Directories
'checkpoint_dir': None,
'sample_dir': None,
'eval_dir': None,
'log_dir': None,
'src_dir': None,
}
if SETUP['model'] == 'bmusegan' and CONFIG['model']['joint_training'] is None:
CONFIG['model']['joint_training'] = SETUP['joint_training']
# Import preset network architectures
if CONFIG['model']['net_g'] is None:
IMPORTED = importlib.import_module(
'.'.join(('musegan', SETUP['model'], 'presets', 'generator',
SETUP['preset_g'])))
CONFIG['model']['net_g'] = IMPORTED.NET_G
if CONFIG['model']['net_d'] is None:
IMPORTED = importlib.import_module(
'.'.join(('musegan', SETUP['model'], 'presets', 'discriminator',
SETUP['preset_d'])))
CONFIG['model']['net_d'] = IMPORTED.NET_D
if SETUP['model'] == 'bmusegan' and CONFIG['model']['net_r'] is None:
IMPORTED = importlib.import_module(
'.'.join(('musegan.bmusegan.presets', 'refiner', SETUP['preset_r'])))
CONFIG['model']['net_r'] = IMPORTED.NET_R
# Set default directories
for kv_pair in (('checkpoint_dir', 'checkpoints'), ('sample_dir', 'samples'),
('eval_dir', 'eval'), ('log_dir', 'logs'), ('src_dir', 'src')):
if CONFIG['model'][kv_pair[0]] is None:
CONFIG['model'][kv_pair[0]] = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'exp', SETUP['model'],
CONFIG['exp']['exp_name'], kv_pair[1])
#===============================================================================
#=================== Make directories & Backup source code =====================
#===============================================================================
# Make sure directories exist
for path in (CONFIG['model']['checkpoint_dir'], CONFIG['model']['sample_dir'],
CONFIG['model']['eval_dir'], CONFIG['model']['log_dir'],
CONFIG['model']['src_dir']):
if not os.path.exists(path):
os.makedirs(path)
# Backup source code
for path in os.listdir(os.path.dirname(os.path.realpath(__file__))):
if os.path.isfile(path):
if path.endswith('.py'):
shutil.copyfile(os.path.basename(path),
os.path.join(CONFIG['model']['src_dir'],
os.path.basename(path)))
distutils.dir_util.copy_tree(
os.path.join(os.path.dirname(os.path.realpath(__file__)), 'musegan'),
os.path.join(CONFIG['model']['src_dir'], 'musegan')
)