tensorflow-trainer

Overview

TensorFlow/Keras model training skill with callbacks, distributed strategies, TensorBoard integration, and production-ready model export capabilities.

Capabilities

• Keras model training with callbacks
• Custom training loops with tf.GradientTape
• Distribution strategy configuration (MirroredStrategy, MultiWorkerMirroredStrategy, TPUStrategy)
• TensorBoard logging and visualization
• SavedModel export for TF Serving
• TFLite conversion for edge deployment
• Mixed precision training

Target Processes

• Model Training Pipeline with Experiment Tracking
• Distributed Training Orchestration
• Model Deployment Pipeline

Tools and Libraries

• TensorFlow
• Keras
• TensorBoard
• TensorFlow Serving
• TensorFlow Lite

Input Schema

{
  "type": "object",
  "required": ["modelConfig", "dataConfig", "trainingConfig"],
  "properties": {
    "modelConfig": {
      "type": "object",
      "properties": {
        "modelPath": { "type": "string" },
        "modelType": { "type": "string", "enum": ["sequential", "functional", "subclassed"] }
      }
    },
    "dataConfig": {
      "type": "object",
      "properties": {
        "trainPath": { "type": "string" },
        "valPath": { "type": "string" },
        "batchSize": { "type": "integer" },
        "prefetch": { "type": "boolean" }
      }
    },
    "trainingConfig": {
      "type": "object",
      "properties": {
        "epochs": { "type": "integer" },
        "optimizer": { "type": "string" },
        "learningRate": { "type": "number" },
        "loss": { "type": "string" },
        "metrics": { "type": "array", "items": { "type": "string" } },
        "callbacks": { "type": "array", "items": { "type": "string" } },
        "distributionStrategy": { "type": "string" }
      }
    },
    "exportConfig": {
      "type": "object",
      "properties": {
        "savedModelPath": { "type": "string" },
        "tflitePath": { "type": "string" },
        "servingSignatures": { "type": "array", "items": { "type": "string" } }
      }
    }
  }
}

Output Schema

{
  "type": "object",
  "required": ["status", "metrics", "modelPath"],
  "properties": {
    "status": {
      "type": "string",
      "enum": ["success", "error", "early_stopped"]
    },
    "metrics": {
      "type": "object",
      "properties": {
        "loss": { "type": "number" },
        "valLoss": { "type": "number" },
        "accuracy": { "type": "number" },
        "valAccuracy": { "type": "number" },
        "epochsTrained": { "type": "integer" }
      }
    },
    "modelPath": {
      "type": "string"
    },
    "savedModelPath": {
      "type": "string"
    },
    "tensorboardLogDir": {
      "type": "string"
    },
    "history": {
      "type": "object",
      "description": "Training history with all metrics per epoch"
    }
  }
}

Usage Example

{
  kind: 'skill',
  title: 'Train TensorFlow model',
  skill: {
    name: 'tensorflow-trainer',
    context: {
      modelConfig: {
        modelPath: 'models/cnn_model.py',
        modelType: 'functional'
      },
      dataConfig: {
        trainPath: 'data/train',
        valPath: 'data/val',
        batchSize: 64,
        prefetch: true
      },
      trainingConfig: {
        epochs: 50,
        optimizer: 'adam',
        learningRate: 0.001,
        loss: 'sparse_categorical_crossentropy',
        metrics: ['accuracy'],
        callbacks: ['early_stopping', 'model_checkpoint', 'tensorboard']
      }
    }
  }
}