Spaces:
Sleeping
Sleeping
| """Train BagOfWords on AG News (12k headlines, 4 topics); about 1 s per epoch on a laptop CPU. | |
| python train.py --epochs 8 | |
| Lines marked `# forge` are the patch the forge setup adds; the rest is a plain training script. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import torch | |
| import torch.nn.functional as F | |
| import forge # forge | |
| from data import CLASSES, HERE, get_loaders | |
| from model import BagOfWords | |
| CHECKPOINT = HERE / "checkpoints" / "best.pt" | |
| BEST = CHECKPOINT.with_suffix(".json") # val acc of best.pt, so a weaker run never replaces it | |
| def saved_best_acc() -> float: | |
| try: | |
| return float(json.loads(BEST.read_text())["val_acc"]) | |
| except (OSError, ValueError, KeyError, TypeError): | |
| return 0.0 | |
| def run_epoch(model, loader, opt) -> tuple[float, float]: | |
| training = opt is not None | |
| model.train(training) | |
| total_loss, correct, n = 0.0, 0, 0 | |
| with torch.set_grad_enabled(training): | |
| for x, y in loader: | |
| out = model(x) | |
| loss = F.cross_entropy(out, y) | |
| if training: | |
| opt.zero_grad() | |
| loss.backward() | |
| opt.step() | |
| total_loss += loss.item() * len(y) | |
| correct += (out.argmax(1) == y).sum().item() | |
| n += len(y) | |
| return total_loss / n, correct / n | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--epochs", type=int, default=8) | |
| parser.add_argument("--data", default=str(HERE / "data")) | |
| parser.add_argument("--batch-size", type=int, default=64) | |
| parser.add_argument("--lr", type=float, default=3e-3) | |
| parser.add_argument("--seed", type=int, default=0) | |
| args = parser.parse_args() | |
| torch.manual_seed(args.seed) | |
| train_loader, val_loader, vocab = get_loaders(args.data, args.batch_size) | |
| model = BagOfWords(len(vocab), len(CLASSES)) # CPU: the model is small enough that a GPU wouldn't help | |
| opt = torch.optim.Adam(model.parameters(), lr=args.lr) | |
| forge.watch(model, val_loader.dataset, classes=CLASSES, epochs=args.epochs) # forge | |
| best_acc = run_best = 0.0 | |
| kept_acc = saved_best_acc() | |
| for epoch in range(1, args.epochs + 1): | |
| train_loss, train_acc = run_epoch(model, train_loader, opt) | |
| val_loss, val_acc = run_epoch(model, val_loader, None) | |
| print(f"epoch {epoch}/{args.epochs} train {train_loss:.3f} / {train_acc:.3f} val {val_loss:.3f} / {val_acc:.3f}") | |
| forge.log(epoch=epoch, train_loss=train_loss, train_acc=train_acc, val_loss=val_loss, # forge | |
| val_acc=val_acc, lr=opt.param_groups[0]["lr"], model=model) # forge | |
| run_best = max(run_best, val_acc) | |
| if val_acc > max(best_acc, kept_acc): | |
| best_acc = val_acc | |
| CHECKPOINT.parent.mkdir(exist_ok=True) | |
| torch.save(model.state_dict(), CHECKPOINT) | |
| BEST.write_text(json.dumps({"val_acc": val_acc, "epoch": epoch, "epochs": args.epochs})) | |
| forge.evaluate(model, val_loader) # forge | |
| if best_acc: | |
| print(f"best val acc {best_acc:.3f}, saved to {CHECKPOINT}") | |
| else: | |
| print(f"best val acc {run_best:.3f}; kept {CHECKPOINT} ({kept_acc:.3f} from an earlier run)") | |
| if __name__ == "__main__": | |
| main() | |