trainModel-tfjscnn.js

async function trainModelFunction(model, data, epochs) {
    const metrics = ['loss', 'val_loss', 'acc', 'val_acc'];
    const container = {
      name: 'Model Training', styles: { height: '1000px' }
    };
    const fitCallbacks = tfvis.show.fitCallbacks(container, metrics);
    
    const batchSize = 512;
  
    const [trainX, trainY] = getBatch(data, 5500);
    const [testX, testY] = getBatch(data, 1000, true);
  
    return model.fit(trainX, trainY, {
      batchSize: batchSize,
      validationData: [testX, testY],
      epochs: epochs,
      shuffle: true,
      callbacks: fitCallbacks
    });
}