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)