| 49 | |
| 50 | // Our actual model |
| 51 | function model(inputXs) { |
| 52 | const xs = inputXs.as4D(-1, IMAGE_SIZE, IMAGE_SIZE, 1); |
| 53 | |
| 54 | const strides = 2; |
| 55 | const pad = 0; |
| 56 | |
| 57 | // Conv 1 |
| 58 | const layer1 = tf.tidy(() => { |
| 59 | return xs.conv2d(conv1Weights, 1, 'same') |
| 60 | .relu() |
| 61 | .maxPool([2, 2], strides, pad); |
| 62 | }); |
| 63 | |
| 64 | // Conv 2 |
| 65 | const layer2 = tf.tidy(() => { |
| 66 | return layer1.conv2d(conv2Weights, 1, 'same') |
| 67 | .relu() |
| 68 | .maxPool([2, 2], strides, pad); |
| 69 | }); |
| 70 | |
| 71 | // Final layer |
| 72 | return layer2.as2D(-1, fullyConnectedWeights.shape[0]) |
| 73 | .matMul(fullyConnectedWeights) |
| 74 | .add(fullyConnectedBias); |
| 75 | } |
| 76 | |
| 77 | // Train the model. |
| 78 | export async function train(data, log) { |