Tensorflow.js, How can I minimize the series-connected models each by each

Question

I'm trying to make a chess-like game AI with 3 ANN models.
when it runs the single minimize method with three models, it makes an error with Failed to compile fragment shader. message.
I'm guessing it's too heavy to handle the three models at once.

These models are all connected. (model1 output->model2 input, model2 output->model3 input)
is there a way to optimize models each by each?
I don't know how to calculate the loss for model1 and model2.


// Partial code

// one minimize contains the 3 model predicts. but I want to split these.
optimizer.minimize(()=>{
    // Model1 predict
    const pieceOutputs = tf.tidy(()=>{
        const inputTensor = tf.tensor4d(makeActionInput(pieceTypes), [pieceTypes.length, 23, 23, 10]);
        return model1.predict(inputTensor).unstack();
    });

    // Model2 predict
    const inputs = makeInputs(batch, pieceOutputs);
    const middles = model2.predict(inputs);
    inputs.dispose(); // trying to call 'dispose' offensively because of the memory issue.

    // Model3 predict
    const extras = makeExtraInputTensor(batch);
    const concated = tf.concat([middles, extras], 1);
    extras.dispose();
    middles.dispose();
    const currentResult = model3.predict(concated);
    concated.dispose();

    // loss
    const targetResult = getTargetResult(batch);
    return tf.mean(currentResult.squaredDifference(targetResult));
});

* Models summary
* Model1 - Piece properties from the movements map of the piece.
Conv Layer x3 (inputShape=[23, 23, 10], kernal=6, filter=15, elu)
Flatten Layer
Dense Layer x3 (units=23*23*10, elu)
Dense Layer (units=2500, elu)
Dense Layer (units=1800, elu)
Dense Layer (units=20, sigmoid)

* Model2 - input = 12*12*(model1.output + extra 4 units)
Conv Layer x3 (inputShape=[12,12,24], kernel=6, filter=30, elu)
Flatten Layer

* Model3 - input = model2.output + extra 2 units
Dense Layer x3 (units=12*12*24+2, elu)
Dense Layer (units=3400, elu)
Dense Layer (units=1200, elu)
Dense Layer (units=81, sigmoid)

Tensorflow.js, How can I minimize the series-connected models each by each

0 Answers0