Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
const [noise, sampledLabels, trick] = tf.tidy(() => {
// Make new latent vectors.
const zVectors = tf.randomUniform([batchSize, latentSize], -1, 1);
const sampledLabels =
tf.randomUniform([batchSize, 1], 0, NUM_CLASSES, 'int32')
.asType('float32');
// We want to train the generator to trick the discriminator.
// For the generator, we want all the {fake, not-fake} labels to say
// not-fake.
const trick = tf.tidy(() => tf.ones([batchSize, 1]).mul(SOFT_ONE));
return [zVectors, sampledLabels, trick];
});
learn(experiences, gamma, tau=TAU) {
tf.tidy(() => {
const tensorified = {};
Object.keys(experiences).map(function(key) {
tensorified[key] = tf.tensor(experiences[key]);
});
const {states, actions, rewards, nextStates, dones} = tensorified;
// Get predicted next-state actions and Q values from target models
const actionsNext = this.actorTarget.predict(nextStates);
const qTargetsNext = this.criticTarget.predict([nextStates, actionsNext]);
// Critic update
this.criticOptimizer.minimize(() => {
// Compute Q targets for current states (y-i)
const qTargets = tf.add(rewards, tf.mul(tf.mul(gamma, qTargetsNext), tf.sub(1, dones)));
const qExpected = this.critic.predict([states, actions]);
const criticLoss = tf.losses.meanSquaredError(qExpected, qTargets);
// torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 1)
async function trainDiscriminatorOneStep(
xTrain, yTrain, batchStart, batchSize, latentSize, generator,
discriminator) {
// TODO(cais): Remove tidy() once the current memory leak issue in tfjs-node
// and tfjs-node-gpu is fixed.
const [x, y, auxY] = tf.tidy(() => {
const imageBatch = xTrain.slice(batchStart, batchSize);
const labelBatch = yTrain.slice(batchStart, batchSize).asType('float32');
// Latent vectors.
let zVectors = tf.randomUniform([batchSize, latentSize], -1, 1);
let sampledLabels =
tf.randomUniform([batchSize, 1], 0, NUM_CLASSES, 'int32')
.asType('float32');
const generatedImages =
generator.predict([zVectors, sampledLabels], {batchSize: batchSize});
const x = tf.concat([imageBatch, generatedImages], 0);
const y = tf.tidy(
() => tf.concat(
async function trainCombinedModelOneStep(batchSize, latentSize, combined) {
// TODO(cais): Remove tidy() once the current memory leak issue in tfjs-node
// and tfjs-node-gpu is fixed.
const [noise, sampledLabels, trick] = tf.tidy(() => {
// Make new latent vectors.
const zVectors = tf.randomUniform([batchSize, latentSize], -1, 1);
const sampledLabels =
tf.randomUniform([batchSize, 1], 0, NUM_CLASSES, 'int32')
.asType('float32');
// We want to train the generator to trick the discriminator.
// For the generator, we want all the {fake, not-fake} labels to say
// not-fake.
const trick = tf.tidy(() => tf.ones([batchSize, 1]).mul(SOFT_ONE));
return [zVectors, sampledLabels, trick];
});
const losses = await combined.trainOnBatch(
[noise, sampledLabels], [trick, sampledLabels]);
tf.dispose([noise, sampledLabels, trick]);
const [x, y, auxY] = tf.tidy(() => {
const imageBatch = xTrain.slice(batchStart, batchSize);
const labelBatch = yTrain.slice(batchStart, batchSize).asType('float32');
// Latent vectors.
let zVectors = tf.randomUniform([batchSize, latentSize], -1, 1);
let sampledLabels =
tf.randomUniform([batchSize, 1], 0, NUM_CLASSES, 'int32')
.asType('float32');
const generatedImages =
generator.predict([zVectors, sampledLabels], {batchSize: batchSize});
const x = tf.concat([imageBatch, generatedImages], 0);
const y = tf.tidy(
() => tf.concat(
[tf.ones([batchSize, 1]).mul(SOFT_ONE), tf.zeros([batchSize, 1])]));
const auxY = tf.concat([labelBatch, sampledLabels], 0);
return [x, y, auxY];
});
async act (state, train=true) {
const action = tf.tidy(() => {
let action = tf.squeeze(this.actor.predict(tf.tensor([state])));
if (train) {
const noise = softmax(this.noise.sample());
action = action.mul(1-this.epsilon).add(tf.mul(noise, this.epsilon));
}
return action;
});
const data = await action.data();
return Array.from(data);
}