Main training loop. {{Describe what this training loop does and which paper sections it follows.}}
(config_path: str = "configs/base.yaml")
| 116 | |
| 117 | |
| 118 | def train(config_path: str = "configs/base.yaml"): |
| 119 | """Main training loop. |
| 120 | |
| 121 | {{Describe what this training loop does and which paper sections it follows.}} |
| 122 | """ |
| 123 | config = load_config(config_path) |
| 124 | |
| 125 | # Device |
| 126 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 127 | |
| 128 | # Build model |
| 129 | model_config = ModelConfig( |
| 130 | # REPLACE: populate from config dict |
| 131 | ) |
| 132 | model = {{MODEL_CLASS}}(model_config).to(device) |
| 133 | |
| 134 | # Build optimizer and scheduler |
| 135 | optimizer = build_optimizer(model, config) |
| 136 | scheduler = build_scheduler(optimizer, config) |
| 137 | |
| 138 | # Build loss |
| 139 | loss_fn = {{LOSS_FN}}() # REPLACE with actual loss construction |
| 140 | |
| 141 | # Build data — see src/data.py for dataset setup instructions |
| 142 | # dataset = {{DATASET_CLASS}}(config["data"]) |
| 143 | # dataloader = DataLoader(dataset, batch_size=config["training"]["batch_size"], |
| 144 | # shuffle=True, num_workers=4, pin_memory=True) |
| 145 | |
| 146 | # Training loop |
| 147 | # §{{SECTION}} — training procedure |
| 148 | gradient_clip = config["training"].get("gradient_clip", None) |
| 149 | total_steps = config["training"]["total_steps"] |
| 150 | |
| 151 | model.train() |
| 152 | step = 0 |
| 153 | |
| 154 | print(f"Starting training for {total_steps} steps") |
| 155 | print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}") |
| 156 | |
| 157 | # REPLACE: implement the actual training loop |
| 158 | # for epoch in range(max_epochs): |
| 159 | # for batch in dataloader: |
| 160 | # batch = {k: v.to(device) for k, v in batch.items()} |
| 161 | # |
| 162 | # optimizer.zero_grad() |
| 163 | # output = model(batch["input"]) |
| 164 | # loss = loss_fn(output, batch["target"]) |
| 165 | # loss.backward() |
| 166 | # |
| 167 | # if gradient_clip is not None: |
| 168 | # torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clip) |
| 169 | # |
| 170 | # optimizer.step() |
| 171 | # if scheduler is not None: |
| 172 | # scheduler.step() |
| 173 | # |
| 174 | # step += 1 |
| 175 | # if step % 100 == 0: |
no test coverage detected