diff --git a/nsnet2/nsnet2.js b/nsnet2/nsnet2.js index 7d41d4c2..3ab0466d 100644 --- a/nsnet2/nsnet2.js +++ b/nsnet2/nsnet2.js @@ -50,7 +50,10 @@ export class NSNet2 { }); const [gru94, gru93] = this.builder_.gru(transpose31, weight192, recurrentWeight193, frames, this.hiddenSize, {bias: bias194, recurrentBias: recurrentBias194, initialHiddenState: initialState92, returnSequence: true}); - const squeeze95 = this.builder_.squeeze(gru93, {axes: [1]}); + // Use reshape to implement squeeze(gru93, {axes: [1]}); + const squeeze95Shape = gru93.shape(); + squeeze95Shape.splice(1, 1); + const squeeze95 = this.builder_.reshape(gru93, squeeze95Shape); const initialState155 = this.builder_.input('initialState155', { type: 'float32', dataType: 'float32', @@ -58,7 +61,10 @@ export class NSNet2 { }); const [gru157, gru156] = this.builder_.gru(squeeze95, weight212, recurrentWeight213, frames, this.hiddenSize, {bias: bias214, recurrentBias: recurrentBias214, initialHiddenState: initialState155, returnSequence: true}); - const squeeze158 = this.builder_.squeeze(gru156, {axes: [1]}); + // Use reshape to implement squeeze(gru156, {axes: [1]}); + const squeeze158Shape = gru156.shape(); + squeeze158Shape.splice(1, 1); + const squeeze158 = this.builder_.reshape(gru156, squeeze158Shape); const transpose159 = this.builder_.transpose(squeeze158, {permutation: [1, 0, 2]}); const relu163 = this.builder_.relu(this.builder_.add(this.builder_.matmul(transpose159, weight215), biasFcOut0)); const relu167 = this.builder_.relu(this.builder_.add(this.builder_.matmul(relu163, weight216), biasFcOut2)); diff --git a/style_transfer/fast_style_transfer_net.js b/style_transfer/fast_style_transfer_net.js index 05e851ac..a2a19f81 100644 --- a/style_transfer/fast_style_transfer_net.js +++ b/style_transfer/fast_style_transfer_net.js @@ -23,8 +23,12 @@ export class FastStyleTransferNet { buildInstanceNormalization_(conv2D, variableMul, variableAdd) { if ('instanceNormalization' in this.builder_) { - return this.builder_.instanceNormalization(conv2D, - {scale: this.builder_.squeeze(variableMul), bias: this.builder_.squeeze(variableAdd)}); + // Use reshape to implement squeeze(variableMul); and squeeze(variableAdd); + const mulShape = variableMul.shape().filter((dim) => dim !==1); + const addShape = variableAdd.shape().filter((dim) => dim !==1); + const mulSqueeze = this.builder_.reshape(variableMul, mulShape); + const addSqueeze = this.builder_.reshape(variableAdd, addShape); + return this.builder_.instanceNormalization(conv2D, {scale: mulSqueeze, bias: addSqueeze}); } else { const sub = this.builder_.sub(conv2D, this.builder_.reduceMean(conv2D, {axes: [2, 3], keepDimensions: true})); const reduceMean = this.builder_.reduceMean(this.builder_.mul(sub, sub), {axes: [2, 3], keepDimensions: true});