def _build_model(self):
acts = dict(
elu=tf.nn.elu,
relu=tf.nn.relu,
swish=tf.nn.swish,
leaky_relu=tf.nn.leaky_relu,
)
cnn_act = acts[self._c.cnn_act]
act = acts[self._c.dense_act]
self._encode = models.ConvEncoder(self._c.cnn_depth, cnn_act)
######################################################################
# RE3: Random Encoder / RunningMeanStd Modules
self._rand_encode = models.ConvRandEncoder(self._c.cnn_depth, cnn_act)
self._rms = models.RMS()
######################################################################
self._dynamics = models.RSSM(
self._c.stoch_size, self._c.deter_size, self._c.deter_size
)
self._decode = models.ConvDecoder(self._c.cnn_depth, cnn_act)
self._reward = models.DenseDecoder((), 2, self._c.num_units, act=act)
if self._c.pcont:
self._pcont = models.DenseDecoder(
(), 3, self._c.num_units, "binary", act=act
)
self._value = models.DenseDecoder((), 3, self._c.num_units, act=act)
self._actor = models.ActionDecoder(
self._actdim,
4,
self._c.num_units,
self._c.action_dist,
init_std=self._c.action_init_std,
act=act,
)
model_modules = [self._encode, self._dynamics, self._decode, self._reward]
if self._c.pcont:
model_modules.append(self._pcont)
Optimizer = functools.partial(
tools.Adam,
wd=self._c.weight_decay,
clip=self._c.grad_clip,
wdpattern=self._c.weight_decay_pattern,
)
self._model_opt = Optimizer("model", model_modules, self._c.model_lr)
self._value_opt = Optimizer("value", [self._value], self._c.value_lr)
self._actor_opt = Optimizer("actor", [self._actor], self._c.actor_lr)
# Do a train step to initialize all variables, including optimizer
# statistics. Ideally, we would use batch size zero, but that doesn't work
# in multi-GPU mode.
self.train(next(self._dataset))
shuffle_data=False,
num_examples=-1,
shuffle_buffer_size=50,
repeat=False)
wav = dataset_train.take(1)
example = iter(dataset_train).next()
mix = example[0][0, :, :].numpy()
sources = example[1][0, :, :].numpy()
write("wavs/mixture.wav", 16000, mix)
write("wavs/source0.wav", 16000, sources[0, :])
write("wavs/source1.wav", 16000, sources[1, :])
write("wavs/source2.wav", 16000, sources[2, :])
write("wavs/source3.wav", 16000, sources[3, :])
mixture = next(iter(dataset_train))[0]
sources = next(iter(dataset_train))[1]
enc = models.ConvEncoder(params=encoder_params,
input_signal_shape=mixture.shape)
enc_output = enc(mixture)
sep = models.TemporalConvNet(separator_params,
input_signal_shape=enc_output.shape)
sep_output = sep(enc_output)
masks = tf.keras.activations.sigmoid(sep_output)
print('bye')
def _build_model(self):
acts = dict(elu=tf.nn.elu,
relu=tf.nn.relu,
swish=tf.nn.swish,
leaky_relu=tf.nn.leaky_relu)
cnn_act = acts[self._c.cnn_act]
act = acts[self._c.dense_act]
#Create encoder based on environment observations
if self._c.proprio:
if self._c.im_size == 64:
self._encode = models.ConvEncoderProprio(
self._c.cnn_depth, cnn_act)
else:
self._encode = models.ConvEncoderProprioLarge(
self._c.cnn_depth, cnn_act)
else:
if self._c.im_size == 64:
self._encode = models.ConvEncoder(self._c.cnn_depth, cnn_act)
else:
self._encode = models.ConvEncoderLarge(self._c.cnn_depth,
cnn_act)
#RSSM model with ensables
self._dynamics = models.RSSME(self._c.stoch_size,
self._c.deter_size,
self._c.deter_size,
num_models=self._c.num_models)
#Create decoder based on image size
if self._c.im_size == 64:
self._decode = models.ConvDecoder(self._c.cnn_depth,
cnn_act,
shape=(self._c.im_size,
self._c.im_size, 3))
else:
self._decode = models.ConvDecoderLarge(self._c.cnn_depth,
cnn_act,
shape=(self._c.im_size,
self._c.im_size, 3))
if self._c.proprio:
self._proprio = models.DenseDecoder((self._propriodim, ),
3,
self._c.num_units,
act=act)
if self._c.pcont:
self._pcont = models.DenseDecoder((),
3,
self._c.num_units,
'binary',
act=act)
self._reward = models.DenseDecoder((), 2, self._c.num_units, act=act)
model_modules = [
self._encode, self._dynamics, self._decode, self._reward
]
if self._c.proprio:
model_modules.append(self._proprio)
if self._c.pcont:
model_modules.append(self._pcont)
#Build actor-critic networks
self._qf1 = models.DenseNetwork(1, 3, self._c.num_units, act=act)
self._qf2 = models.DenseNetwork(1, 3, self._c.num_units, act=act)
self._target_qf1 = deepcopy(self._qf2)
self._target_qf2 = deepcopy(self._qf1)
self._qf_criterion = tf.keras.losses.Huber()
self._actor = models.ActorNetwork(self._actdim,
4,
self._c.num_units,
act=act)
#Initialize optimizers
Optimizer = functools.partial(tools.Adam,
wd=self._c.weight_decay,
clip=self._c.grad_clip,
wdpattern=self._c.weight_decay_pattern)
self._model_opt = Optimizer('model', model_modules, self._c.model_lr)
self._qf_opt = Optimizer('qf', [self._qf1, self._qf2], self._c.q_lr)
self._actor_opt = Optimizer('actor', [self._actor], self._c.actor_lr)
