Agent skill
pytorch-trainer
PyTorch model training skill with custom training loops, gradient management, and GPU optimization.
Install this agent skill to your Project
npx add-skill https://github.com/a5c-ai/babysitter/tree/main/library/specializations/data-science-ml/skills/pytorch-trainer
SKILL.md
pytorch-trainer
Overview
PyTorch model training skill with custom training loops, gradient management, GPU optimization, and integration with experiment tracking systems.
Capabilities
- Custom training loop execution
- Learning rate scheduling (StepLR, CosineAnnealing, OneCycleLR, etc.)
- Gradient clipping and accumulation
- Mixed precision training (AMP)
- Checkpoint management and resumption
- DataLoader optimization
- Multi-GPU training (DataParallel, DistributedDataParallel)
- Early stopping with patience
Target Processes
- Model Training Pipeline with Experiment Tracking
- Distributed Training Orchestration
- AutoML Pipeline Orchestration
Tools and Libraries
- PyTorch
- PyTorch Lightning (optional)
- torchvision, torchaudio, torchtext
- CUDA toolkit
Input Schema
{
"type": "object",
"required": ["modelPath", "dataConfig", "trainingConfig"],
"properties": {
"modelPath": {
"type": "string",
"description": "Path to model definition file"
},
"dataConfig": {
"type": "object",
"properties": {
"trainPath": { "type": "string" },
"valPath": { "type": "string" },
"batchSize": { "type": "integer" },
"numWorkers": { "type": "integer" }
}
},
"trainingConfig": {
"type": "object",
"properties": {
"epochs": { "type": "integer" },
"learningRate": { "type": "number" },
"optimizer": { "type": "string" },
"scheduler": { "type": "string" },
"mixedPrecision": { "type": "boolean" },
"gradientClipping": { "type": "number" },
"gradientAccumulation": { "type": "integer" }
}
},
"checkpointConfig": {
"type": "object",
"properties": {
"saveDir": { "type": "string" },
"saveEvery": { "type": "integer" },
"resumeFrom": { "type": "string" }
}
}
}
}
Output Schema
{
"type": "object",
"required": ["status", "metrics", "checkpointPath"],
"properties": {
"status": {
"type": "string",
"enum": ["success", "error", "early_stopped"]
},
"metrics": {
"type": "object",
"properties": {
"trainLoss": { "type": "number" },
"valLoss": { "type": "number" },
"trainAccuracy": { "type": "number" },
"valAccuracy": { "type": "number" },
"epochsTrained": { "type": "integer" },
"trainingTime": { "type": "number" }
}
},
"checkpointPath": {
"type": "string"
},
"learningCurve": {
"type": "array",
"items": {
"type": "object",
"properties": {
"epoch": { "type": "integer" },
"trainLoss": { "type": "number" },
"valLoss": { "type": "number" }
}
}
}
}
}
Usage Example
{
kind: 'skill',
title: 'Train PyTorch model',
skill: {
name: 'pytorch-trainer',
context: {
modelPath: 'models/resnet.py',
dataConfig: {
trainPath: 'data/train',
valPath: 'data/val',
batchSize: 32,
numWorkers: 4
},
trainingConfig: {
epochs: 100,
learningRate: 0.001,
optimizer: 'AdamW',
scheduler: 'cosine',
mixedPrecision: true,
gradientClipping: 1.0
}
}
}
}
Recommended Agent Skills
Expand your agent's capabilities with these related and highly-rated skills.
gsd-tools
Central utility skill for GSD operations. Provides config parsing, slug generation, timestamps, path operations, and orchestrates calls to other specialized skills. Acts as the unified entry point that the original gsd-tools.cjs provided via its lib/ modules (commands, config, core, init).
model-profile-resolution
Resolve model profile (quality/balanced/budget) at orchestration start and map agents to specific models. Enables cost/quality tradeoffs by selecting appropriate AI models for each agent role.
verification-suite
Plan structure validation, phase completeness checks, reference integrity verification, and artifact existence confirmation. Provides the structured verification layer ensuring GSD artifacts are well-formed and complete.
state-management
STATE.md reading, writing, and field-level updates. Provides cross-session state persistence via .planning/STATE.md with structured fields for current task, completed phases, blockers, decisions, and quick tasks.
git-integration
Git commit patterns, formats, and conventions for GSD methodology. Provides atomic commits per task, structured commit messages, planning file commits, branch management, and milestone tag operations.
frontmatter-parsing
YAML frontmatter parsing and manipulation for .planning/ documents. Provides read, write, update, query, and validation operations on frontmatter blocks in GSD markdown artifacts.
Didn't find tool you were looking for?