{"id":15418409,"url":"https://github.com/idorce/torchtrain","last_synced_at":"2025-10-11T08:31:03.964Z","repository":{"id":57476577,"uuid":"217583462","full_name":"idorce/torchtrain","owner":"idorce","description":"A small tool for PyTorch training","archived":false,"fork":false,"pushed_at":"2021-03-14T13:39:43.000Z","size":125,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"master","last_synced_at":"2025-01-06T14:54:43.058Z","etag":null,"topics":["deep-learning","machine-learning","ml","neural-network","python","pytorch","pytorch-tool","pytorch-training"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/idorce.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2019-10-25T17:28:31.000Z","updated_at":"2021-03-14T13:38:53.000Z","dependencies_parsed_at":"2022-09-12T14:50:16.765Z","dependency_job_id":null,"html_url":"https://github.com/idorce/torchtrain","commit_stats":null,"previous_names":[],"tags_count":42,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/idorce%2Ftorchtrain","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/idorce%2Ftorchtrain/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/idorce%2Ftorchtrain/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/idorce%2Ftorchtrain/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/idorce","download_url":"https://codeload.github.com/idorce/torchtrain/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":236067846,"owners_count":19089796,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","machine-learning","ml","neural-network","python","pytorch","pytorch-tool","pytorch-training"],"created_at":"2024-10-01T17:20:40.008Z","updated_at":"2025-10-11T08:30:58.657Z","avatar_url":"https://github.com/idorce.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# 🔥 torchtrain 💪\n\nA small tool for PyTorch training.\n\n## Features\n\n- Avoid boilerplate code for training.\n- Stepwise training.\n- Automatic TensorBoard logging, and tqdm bar.\n- Count model parameters and save hyperparameters.\n- DataParallel.\n- Early stop.\n- Save and load checkpoint. Continue training.\n- Catch out of memory exceptions to avoid breaking training.\n- Gradient accumulation.\n- Gradient clipping.\n- Only run few epochs, steps and batches for code test.\n\n## Install\n\n```\npip install torchtrain\n```\n\n## Example\n\nCheck doc string of [`Trainer` class](https://github.com/idorce/torchtrain/blob/master/torchtrain/trainer.py) for detailed configurations.\n\nAn incomplete minimal example:\n\n```python\ndata_iter = get_data()\nmodel = Bert()\noptimizer = Adam(model.parameters(), lr=cfg[\"lr\"])\ncriteria = {\"loss\": AverageAggregator(BCELoss())}\ntrainer = Trainer(model, data_iter, criteria, cfg, optimizer)\ntrainer.train(stepwise=True)\n```\n\nOr this version:\n\n```python\nfrom argparse import ArgumentParser\n\nfrom sklearn.model_selection import ParameterGrid\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom transformers import AutoModel, BertTokenizer\n\nfrom data.load import get_batch_size, get_data\nfrom metrics import BCELoss\nfrom models import BertSumExt\nfrom torchtrain import Trainer\nfrom torchtrain.metrics import AverageAggregator\nfrom torchtrain.utils import set_random_seeds\n\n\ndef get_args():\n    parser = ArgumentParser()\n    parser.add_argument(\"--seed\", type=int, default=233666)\n    parser.add_argument(\"--run_ckp\", default=\"\")\n    parser.add_argument(\"--run_dataset\", default=\"val\")\n    parser.add_argument(\"--batch_size\", type=int, default=64)\n    parser.add_argument(\"--warmup\", type=int, default=10000)\n    parser.add_argument(\"--stepwise\", action=\"store_false\")\n    # torchtrain cfgs\n    parser.add_argument(\"--max_n\", type=int, default=50000)\n    parser.add_argument(\"--val_step\", type=int, default=1000)\n    parser.add_argument(\"--save_path\", default=\"/tmp/runs\")\n    parser.add_argument(\"--model_name\", default=\"BertSumExt\")\n    parser.add_argument(\"--cuda_list\", default=\"2,3\")\n    parser.add_argument(\"--grad_accum_batch\", type=int, default=1)\n    parser.add_argument(\"--train_few\", action=\"store_true\")\n    return vars(parser.parse_args())\n\n\ndef get_param_grid():\n    param_grid = [\n        {\"pretrained_model_name\": [\"voidful/albert_chinese_tiny\"], \"lr\": [6e-5]},\n    ]\n    return ParameterGrid(param_grid)\n\n\ndef get_cfg(args={}, params={}):\n    cfg = {**args, **params}\n    # other cfgs\n    return cfg\n\n\ndef run(cfg):\n    set_random_seeds(cfg[\"seed\"])\n    tokenizer = BertTokenizer.from_pretrained(cfg[\"pretrained_model_name\"])\n    bert = AutoModel.from_pretrained(cfg[\"pretrained_model_name\"])\n    data_iter = get_data(\n        cfg[\"batch_size\"], tokenizer, bert.config.max_position_embeddings\n    )\n    model = BertSumExt(bert)\n    optimizer = Adam(model.parameters(), lr=cfg[\"lr\"])\n    scheduler = LambdaLR(\n        optimizer,\n        lambda step: min(step ** (-0.5), step * (cfg[\"warmup\"] ** (-1.5)))\n        if step \u003e 0\n        else 0,\n    )\n    criteria = {\"loss\": AverageAggregator(BCELoss())}\n    trainer = Trainer(\n        model,\n        data_iter,\n        criteria,\n        cfg,\n        optimizer,\n        scheduler,\n        get_batch_size=get_batch_size,\n    )\n    if cfg[\"run_ckp\"]:\n        return trainer.test(cfg[\"run_ckp\"], cfg[\"run_dataset\"])\n    return trainer.train(stepwise=cfg[\"stepwise\"])\n\n\ndef main():\n    param_grid = get_param_grid()\n    for i, params in enumerate(param_grid):\n        print(\"Config\", str(i + 1), \"/\", str(len(param_grid)))\n        cfg = get_cfg(get_args(), params)\n        metrics = run(cfg)\n        print(\"Best metrics:\", metrics)\n\n\nif __name__ == \"__main__\":\n    main()\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fidorce%2Ftorchtrain","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fidorce%2Ftorchtrain","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fidorce%2Ftorchtrain/lists"}