Lucy: 71M-param unconditional 128x128 face DDPM, trained from scratch

.gitattributes

@@ -33,3 +33,24 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/sample_grid.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0003000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0006000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0009000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0012000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0015000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0018000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0021000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0024000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0027000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0030000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0033000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0036000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0039000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0042000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0045000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0048000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0051000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0054000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0057000.png filter=lfs diff=lfs merge=lfs -text
+training/progress/step_0060000.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,244 @@
+---
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - ddpm
+  - unconditional-image-generation
+  - diffusion
+  - faces
+  - pytorch
+  - safetensors
+datasets:
+  - tglcourse/CelebA-faces-cropped-128
+---
+
+# Lucy
+
+**Lucy** is a small unconditional denoising diffusion model (DDPM) that generates
+128×128 RGB face images. It was trained from scratch — no pretrained
+initialization, no fine-tuning of an existing checkpoint — on a single consumer
+GPU (RTX 3060, 12 GB) in under ten hours.
+
+Weights are open: Apache-2.0, `safetensors`, plus the complete training script
+and the raw training log used to produce them. Everything needed to reproduce
+the run is in this repo.
+
+![Samples from Lucy](samples/sample_grid.png)
+
+---
+
+## Quick start
+
+```bash
+pip install diffusers transformers accelerate torch safetensors
+```
+
+```python
+from diffusers import DDPMPipeline
+
+pipe = DDPMPipeline.from_pretrained("opencerebral/Lucy").to("cuda")
+images = pipe(batch_size=4).images
+images[0].save("face.png")
+```
+
+### Faster sampling with DDIM
+
+The default DDPM scheduler runs the full 1000-step reverse chain. The samples in
+this card were generated with **DDIM at 50 steps**, which is ~20× faster and what
+the training script itself uses for its periodic previews:
+
+```python
+import torch
+from diffusers import DDIMPipeline, DDIMScheduler, UNet2DModel
+
+unet = UNet2DModel.from_pretrained("opencerebral/Lucy", subfolder="unet").to("cuda")
+sched = DDIMScheduler(num_train_timesteps=1000, beta_schedule="linear")
+pipe = DDIMPipeline(unet=unet, scheduler=sched).to("cuda")
+
+images = pipe(batch_size=16, num_inference_steps=50,
+              generator=torch.manual_seed(42)).images
+```
+
+The model is unconditional: there is no text prompt, no class label, and no
+guidance. The only control you have over the output is the initial noise, i.e.
+the random seed.
+
+---
+
+## Model details
+
+| | |
+|---|---|
+| Architecture | `UNet2DModel` (diffusers), epsilon-prediction |
+| Parameters | **71.4 M** |
+| Resolution | 128 × 128, RGB |
+| Conditioning | none (unconditional) |
+| Block channels | `[128, 128, 256, 256, 512]` |
+| Down blocks | `DownBlock2D ×3`, `AttnDownBlock2D`, `DownBlock2D` |
+| Up blocks | `UpBlock2D`, `AttnUpBlock2D`, `UpBlock2D ×3` |
+| Layers per block | 2 |
+| Attention head dim | 8 (self-attention at the 16×16 stage and in the mid block) |
+| Norm | GroupNorm, 32 groups, eps 1e-5 |
+| Activation | SiLU |
+| Timestep embedding | positional |
+| Weights format | `safetensors`, fp32 (286 MB) |
+| diffusers version | 0.39.0 |
+
+### Noise schedule
+
+Linear beta schedule, `beta_start = 1e-4`, `beta_end = 0.02`, `T = 1000`
+timesteps, `prediction_type = "epsilon"`, `variance_type = "fixed_small"`,
+`clip_sample = True`. Standard DDPM (Ho et al., 2020) settings.
+
+Attention sits at exactly one resolution (the second-to-last downsampling
+stage, 16×16). That is deliberate: at 128×128 a full-resolution attention map is
+what blows up memory on a 12 GB card, and a single mid-depth attention stage
+recovers most of the global coherence — face symmetry, consistent lighting
+across the image — for very little cost.
+
+---
+
+## Training
+
+| | |
+|---|---|
+| Dataset | [`tglcourse/CelebA-faces-cropped-128`](https://huggingface.co/datasets/tglcourse/CelebA-faces-cropped-128) — 192,469 rows used |
+| Hardware | 1 × NVIDIA GeForce RTX 3060 (12 GB) |
+| Precision | fp16 AMP (`torch.autocast`) with `GradScaler`; TF32 matmul enabled |
+| Optimizer | AdamW, lr `1e-4`, 500-step linear warmup, constant thereafter |
+| Batch size | 16 (no gradient accumulation) |
+| Grad clipping | global norm 1.0 |
+| EMA | decay 0.9999, updated every optimizer step; **the released weights are the EMA weights** |
+| Steps | 60,000 (≈ 5 epochs) |
+| Wall clock | **9.54 hours**, ~1.75 it/s steady state |
+| Seed | 42 |
+| Loss | MSE on predicted noise, computed in fp32 |
+
+Augmentation was resize → center crop to 128 → random horizontal flip →
+normalize to `[-1, 1]`. Nothing else.
+
+### Loss curve
+
+Training loss over the run, read from `training/train.log`:
+
+| step | loss |
+|---:|---:|
+| 100 | 0.5997 |
+| 200 | 0.0715 |
+| 300 | 0.0467 |
+| 60,000 | 0.0131 |
+
+Note that the *first* logged value is the only large one — DDPM epsilon-loss
+collapses almost immediately to the ~0.05 range because most sampled timesteps
+are easy, and the remaining 59,700 steps buy a slow grind from 0.047 down to
+0.013. Loss is a poor progress signal here; the sample grids are the real one.
+
+### Training progression
+
+`training/progress/` contains an EMA sample grid rendered every 3,000 steps
+(`step_0003000.png` … `step_0060000.png`), all from the same fixed seed, so you
+can watch the same 16 latents resolve from noise into faces over the run.
+
+---
+
+## What's in this repo
+
+```
+model_index.json                 DDPMPipeline definition
+unet/                            UNet2DModel config + EMA weights (safetensors)
+scheduler/                       DDPMScheduler config
+samples/sample_grid.png          16 samples, DDIM 50 steps, seed 42
+samples/one_face.png             single 128x128 sample
+training/train.py                the full training script
+training/run.sh                  the exact launch command for this run
+training/train.log               complete stdout of the 60k-step run
+training/progress/               EMA previews every 3,000 steps
+```
+
+### Reproducing
+
+```bash
+pip install torch torchvision diffusers datasets accelerate tensorboard
+bash training/run.sh          # edit --output_dir first
+```
+
+The script stops at `min(--max_train_steps, --max_hours)`, checkpoints every
+20,000 steps keeping the last 2, and writes a ready-to-load `DDPMPipeline` to
+`<output_dir>/final` when it finishes. It supports `--resume_from` for restarts.
+
+Optimizer-state checkpoints (`step_*.pt`, ~1.2 GB each) are **not** included
+here — they are large and only useful for resuming this specific run. The EMA
+weights in `unet/` are the artifact worth keeping.
+
+---
+
+## Limitations
+
+Be realistic about what a 71M-parameter model trained for 9.5 hours produces:
+
+- **Fidelity.** Outputs are recognizably faces with plausible structure,
+  lighting, and hair, but they are soft, and fine detail — teeth, eyes at close
+  range, jewelry, text — is often wrong. Compare the samples above before
+  assuming this is drop-in usable.
+- **Backgrounds are mush.** CelebA is celebrity photos with wildly varied,
+  cluttered backgrounds. The model has learned "a colorful blur behind a face"
+  and little more.
+- **Occasional structural failures.** A minority of samples have distorted or
+  doubled features. There is no rejection sampling or guidance to filter these.
+- **No control.** Unconditional means unconditional. You cannot ask for an age,
+  an expression, a pose, or a style. Seed search is the only lever.
+- **Dataset bias.** CelebA is a celebrity face dataset and is well documented as
+  skewed in age, skin tone, and attractiveness-as-labeled. Lucy reproduces that
+  skew faithfully. It is not a representative sample of human faces and should
+  not be treated as one.
+- **Resolution.** Fixed at 128×128. The UNet's `sample_size` is 128 and it will
+  degrade badly if you ask it for other sizes.
+
+## Intended use
+
+Research, education, and demonstration of small-scale from-scratch diffusion
+training. It is a good teaching artifact: the whole pipeline is one 270-line
+script, and the log shows exactly what a 60k-step run costs on hardware people
+actually own.
+
+**Not** intended for generating imagery of, or presented as, real people. The
+model was trained on photographs of real individuals; while its outputs are not
+reconstructions of specific training images, do not use it to imply the depiction
+of any real person, and do not use it in any identity, verification, or
+surveillance context.
+
+## Licensing
+
+- **Model weights and code in this repo: Apache-2.0.**
+- **Training data is a separate matter.** The weights derive from CelebA, which
+  is released for **non-commercial research use only** under its original terms.
+  The Apache-2.0 grant here covers what the repo authors can grant; it does not
+  and cannot override upstream dataset terms. If you plan to use Lucy
+  commercially, review the CelebA license yourself first — the honest reading is
+  that a CelebA-derived model inherits that non-commercial restriction.
+
+## Citation
+
+DDPM, the method this implements:
+
+```bibtex
+@inproceedings{ho2020ddpm,
+  title  = {Denoising Diffusion Probabilistic Models},
+  author = {Ho, Jonathan and Jain, Ajay and Abbeel, Pieter},
+  booktitle = {Advances in Neural Information Processing Systems},
+  year   = {2020}
+}
+```
+
+This model:
+
+```bibtex
+@misc{lucy2026,
+  title  = {Lucy: a 71M-parameter unconditional face diffusion model trained from scratch},
+  author = {OpenCerebral},
+  year   = {2026},
+  howpublished = {\url{https://huggingface.co/opencerebral/Lucy}}
+}
+```

model_index.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "DDPMPipeline",
+  "_diffusers_version": "0.39.0",
+  "scheduler": [
+    "diffusers",
+    "DDPMScheduler"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DModel"
+  ]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "DDPMScheduler",
+  "_diffusers_version": "0.39.0",
+  "beta_end": 0.02,
+  "beta_schedule": "linear",
+  "beta_start": 0.0001,
+  "clip_sample": true,
+  "clip_sample_range": 1.0,
+  "dynamic_thresholding_ratio": 0.995,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "steps_offset": 0,
+  "thresholding": false,
+  "timestep_spacing": "leading",
+  "trained_betas": null,
+  "variance_type": "fixed_small"
+}

training/run.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+# Overnight pretraining run (RTX 3060). Launch detached:
+#   nohup bash run.sh > train.log 2>&1 &
+# or inside tmux:  tmux new -s train 'bash run.sh'
+set -euo pipefail
+cd "$(dirname "$0")"
+export HF_HOME=/media/joseph/Media/.hf-cache
+export TOKENIZERS_PARALLELISM=false
+
+exec ./.venv/bin/python -u train.py \
+  --dataset_name tglcourse/CelebA-faces-cropped-128 \
+  --resolution 128 \
+  --train_batch_size 16 \
+  --max_train_steps 60000 \
+  --max_hours 13.5 \
+  --output_dir /media/joseph/Media/diffusion128/runs/celeba128

training/train.log

@@ -0,0 +1,603 @@
+Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
+model params: 71.4M | dataset rows: 192469 | batch 16 x accum 1
+step 100/60000 | epoch 0 | loss 0.5997 | grad 0.28 | 1.49 it/s | ETA 11.16h
+step 200/60000 | epoch 0 | loss 0.0715 | grad 0.17 | 1.56 it/s | ETA 10.64h
+step 300/60000 | epoch 0 | loss 0.0467 | grad 0.42 | 1.60 it/s | ETA 10.33h
+step 400/60000 | epoch 0 | loss 0.0403 | grad 0.37 | 1.63 it/s | ETA 10.15h
+step 500/60000 | epoch 0 | loss 0.0346 | grad 0.59 | 1.65 it/s | ETA 10.02h
+step 600/60000 | epoch 0 | loss 0.0279 | grad 0.25 | 1.66 it/s | ETA 9.93h
+step 700/60000 | epoch 0 | loss 0.0225 | grad 0.17 | 1.67 it/s | ETA 9.86h
+step 800/60000 | epoch 0 | loss 0.0265 | grad 0.15 | 1.68 it/s | ETA 9.80h
+step 900/60000 | epoch 0 | loss 0.0232 | grad 0.11 | 1.68 it/s | ETA 9.76h
+step 1000/60000 | epoch 0 | loss 0.0219 | grad 0.26 | 1.68 it/s | ETA 9.74h
+step 1100/60000 | epoch 0 | loss 0.0199 | grad 0.37 | 1.68 it/s | ETA 9.72h
+step 1200/60000 | epoch 0 | loss 0.0201 | grad 0.06 | 1.68 it/s | ETA 9.72h
+step 1300/60000 | epoch 0 | loss 0.0205 | grad 0.29 | 1.68 it/s | ETA 9.72h
+step 1400/60000 | epoch 0 | loss 0.0195 | grad 0.25 | 1.68 it/s | ETA 9.71h
+step 1500/60000 | epoch 0 | loss 0.0192 | grad 0.20 | 1.68 it/s | ETA 9.70h
+step 1600/60000 | epoch 0 | loss 0.0172 | grad 0.17 | 1.67 it/s | ETA 9.69h
+step 1700/60000 | epoch 0 | loss 0.0173 | grad 0.21 | 1.67 it/s | ETA 9.69h
+step 1800/60000 | epoch 0 | loss 0.0187 | grad 0.18 | 1.67 it/s | ETA 9.68h
+step 1900/60000 | epoch 0 | loss 0.0169 | grad 0.18 | 1.67 it/s | ETA 9.66h
+step 2000/60000 | epoch 0 | loss 0.0166 | grad 0.20 | 1.67 it/s | ETA 9.63h
+step 2100/60000 | epoch 0 | loss 0.0191 | grad 0.10 | 1.67 it/s | ETA 9.60h
+step 2200/60000 | epoch 0 | loss 0.0185 | grad 0.16 | 1.68 it/s | ETA 9.58h
+step 2300/60000 | epoch 0 | loss 0.0191 | grad 0.11 | 1.67 it/s | ETA 9.58h
+step 2400/60000 | epoch 0 | loss 0.0169 | grad 0.16 | 1.67 it/s | ETA 9.59h
+step 2500/60000 | epoch 0 | loss 0.0190 | grad 0.11 | 1.66 it/s | ETA 9.59h
+step 2600/60000 | epoch 0 | loss 0.0195 | grad 0.12 | 1.66 it/s | ETA 9.60h
+step 2700/60000 | epoch 0 | loss 0.0190 | grad 0.07 | 1.66 it/s | ETA 9.60h
+step 2800/60000 | epoch 0 | loss 0.0172 | grad 0.15 | 1.66 it/s | ETA 9.60h
+step 2900/60000 | epoch 0 | loss 0.0154 | grad 0.16 | 1.65 it/s | ETA 9.59h
+step 3000/60000 | epoch 0 | loss 0.0188 | grad 0.12 | 1.65 it/s | ETA 9.59h
+step 3100/60000 | epoch 0 | loss 0.0165 | grad 0.05 | 1.64 it/s | ETA 9.64h
+step 3200/60000 | epoch 0 | loss 0.0166 | grad 0.08 | 1.64 it/s | ETA 9.64h
+step 3300/60000 | epoch 0 | loss 0.0158 | grad 0.13 | 1.64 it/s | ETA 9.63h
+step 3400/60000 | epoch 0 | loss 0.0167 | grad 0.08 | 1.63 it/s | ETA 9.62h
+step 3500/60000 | epoch 0 | loss 0.0165 | grad 0.28 | 1.63 it/s | ETA 9.61h
+step 3600/60000 | epoch 0 | loss 0.0154 | grad 0.25 | 1.63 it/s | ETA 9.61h
+step 3700/60000 | epoch 0 | loss 0.0170 | grad 0.23 | 1.63 it/s | ETA 9.60h
+step 3800/60000 | epoch 0 | loss 0.0168 | grad 0.15 | 1.63 it/s | ETA 9.58h
+step 3900/60000 | epoch 0 | loss 0.0167 | grad 0.08 | 1.63 it/s | ETA 9.54h
+step 4000/60000 | epoch 0 | loss 0.0160 | grad 0.14 | 1.64 it/s | ETA 9.50h
+step 4100/60000 | epoch 0 | loss 0.0178 | grad 0.05 | 1.64 it/s | ETA 9.46h
+step 4200/60000 | epoch 0 | loss 0.0178 | grad 0.17 | 1.65 it/s | ETA 9.42h
+step 4300/60000 | epoch 0 | loss 0.0170 | grad 0.05 | 1.65 it/s | ETA 9.38h
+step 4400/60000 | epoch 0 | loss 0.0164 | grad 0.03 | 1.65 it/s | ETA 9.34h
+step 4500/60000 | epoch 0 | loss 0.0164 | grad 0.18 | 1.66 it/s | ETA 9.31h
+step 4600/60000 | epoch 0 | loss 0.0167 | grad 0.17 | 1.66 it/s | ETA 9.27h
+step 4700/60000 | epoch 0 | loss 0.0160 | grad 0.12 | 1.66 it/s | ETA 9.24h
+step 4800/60000 | epoch 0 | loss 0.0141 | grad 0.07 | 1.67 it/s | ETA 9.20h
+step 4900/60000 | epoch 0 | loss 0.0182 | grad 0.17 | 1.67 it/s | ETA 9.17h
+step 5000/60000 | epoch 0 | loss 0.0154 | grad 0.05 | 1.67 it/s | ETA 9.14h
+step 5100/60000 | epoch 0 | loss 0.0180 | grad 0.09 | 1.67 it/s | ETA 9.10h
+step 5200/60000 | epoch 0 | loss 0.0159 | grad 0.13 | 1.68 it/s | ETA 9.07h
+step 5300/60000 | epoch 0 | loss 0.0152 | grad 0.21 | 1.68 it/s | ETA 9.04h
+step 5400/60000 | epoch 0 | loss 0.0161 | grad 0.21 | 1.68 it/s | ETA 9.01h
+step 5500/60000 | epoch 0 | loss 0.0159 | grad 0.10 | 1.69 it/s | ETA 8.98h
+step 5600/60000 | epoch 0 | loss 0.0153 | grad 0.04 | 1.69 it/s | ETA 8.96h
+step 5700/60000 | epoch 0 | loss 0.0150 | grad 0.13 | 1.69 it/s | ETA 8.93h
+step 5800/60000 | epoch 0 | loss 0.0169 | grad 0.07 | 1.69 it/s | ETA 8.91h
+step 5900/60000 | epoch 0 | loss 0.0169 | grad 0.19 | 1.69 it/s | ETA 8.91h
+step 6000/60000 | epoch 0 | loss 0.0144 | grad 0.16 | 1.68 it/s | ETA 8.91h
+step 6100/60000 | epoch 0 | loss 0.0161 | grad 0.11 | 1.68 it/s | ETA 8.93h
+step 6200/60000 | epoch 0 | loss 0.0156 | grad 0.23 | 1.67 it/s | ETA 8.93h
+step 6300/60000 | epoch 0 | loss 0.0157 | grad 0.12 | 1.67 it/s | ETA 8.93h
+step 6400/60000 | epoch 0 | loss 0.0185 | grad 0.11 | 1.67 it/s | ETA 8.92h
+step 6500/60000 | epoch 0 | loss 0.0141 | grad 0.12 | 1.67 it/s | ETA 8.92h
+step 6600/60000 | epoch 0 | loss 0.0165 | grad 0.10 | 1.66 it/s | ETA 8.91h
+step 6700/60000 | epoch 0 | loss 0.0137 | grad 0.11 | 1.66 it/s | ETA 8.91h
+step 6800/60000 | epoch 0 | loss 0.0178 | grad 0.08 | 1.66 it/s | ETA 8.90h
+step 6900/60000 | epoch 0 | loss 0.0151 | grad 0.12 | 1.66 it/s | ETA 8.90h
+step 7000/60000 | epoch 0 | loss 0.0151 | grad 0.13 | 1.65 it/s | ETA 8.90h
+step 7100/60000 | epoch 0 | loss 0.0164 | grad 0.10 | 1.65 it/s | ETA 8.89h
+step 7200/60000 | epoch 0 | loss 0.0175 | grad 0.11 | 1.65 it/s | ETA 8.88h
+step 7300/60000 | epoch 0 | loss 0.0169 | grad 0.05 | 1.65 it/s | ETA 8.88h
+step 7400/60000 | epoch 0 | loss 0.0165 | grad 0.12 | 1.65 it/s | ETA 8.88h
+step 7500/60000 | epoch 0 | loss 0.0158 | grad 0.10 | 1.64 it/s | ETA 8.88h
+step 7600/60000 | epoch 0 | loss 0.0173 | grad 0.03 | 1.64 it/s | ETA 8.87h
+step 7700/60000 | epoch 0 | loss 0.0134 | grad 0.07 | 1.64 it/s | ETA 8.87h
+step 7800/60000 | epoch 0 | loss 0.0142 | grad 0.11 | 1.64 it/s | ETA 8.86h
+step 7900/60000 | epoch 0 | loss 0.0153 | grad 0.16 | 1.63 it/s | ETA 8.86h
+step 8000/60000 | epoch 0 | loss 0.0159 | grad 0.03 | 1.63 it/s | ETA 8.85h
+step 8100/60000 | epoch 0 | loss 0.0163 | grad 0.05 | 1.63 it/s | ETA 8.85h
+step 8200/60000 | epoch 0 | loss 0.0143 | grad 0.12 | 1.63 it/s | ETA 8.84h
+step 8300/60000 | epoch 0 | loss 0.0157 | grad 0.06 | 1.63 it/s | ETA 8.84h
+step 8400/60000 | epoch 0 | loss 0.0149 | grad 0.09 | 1.62 it/s | ETA 8.83h
+step 8500/60000 | epoch 0 | loss 0.0141 | grad 0.10 | 1.62 it/s | ETA 8.82h
+step 8600/60000 | epoch 0 | loss 0.0149 | grad 0.12 | 1.62 it/s | ETA 8.81h
+step 8700/60000 | epoch 0 | loss 0.0164 | grad 0.07 | 1.62 it/s | ETA 8.81h
+step 8800/60000 | epoch 0 | loss 0.0160 | grad 0.16 | 1.62 it/s | ETA 8.80h
+step 8900/60000 | epoch 0 | loss 0.0151 | grad 0.14 | 1.62 it/s | ETA 8.79h
+step 9000/60000 | epoch 0 | loss 0.0182 | grad 0.08 | 1.61 it/s | ETA 8.78h
+step 9100/60000 | epoch 0 | loss 0.0149 | grad 0.12 | 1.61 it/s | ETA 8.79h
+step 9200/60000 | epoch 0 | loss 0.0173 | grad 0.08 | 1.61 it/s | ETA 8.78h
+step 9300/60000 | epoch 0 | loss 0.0143 | grad 0.10 | 1.61 it/s | ETA 8.77h
+step 9400/60000 | epoch 0 | loss 0.0135 | grad 0.12 | 1.60 it/s | ETA 8.76h
+step 9500/60000 | epoch 0 | loss 0.0152 | grad 0.08 | 1.61 it/s | ETA 8.74h
+step 9600/60000 | epoch 0 | loss 0.0141 | grad 0.09 | 1.61 it/s | ETA 8.71h
+step 9700/60000 | epoch 0 | loss 0.0159 | grad 0.09 | 1.61 it/s | ETA 8.69h
+step 9800/60000 | epoch 0 | loss 0.0161 | grad 0.08 | 1.61 it/s | ETA 8.66h
+step 9900/60000 | epoch 0 | loss 0.0156 | grad 0.08 | 1.61 it/s | ETA 8.64h
+step 10000/60000 | epoch 0 | loss 0.0136 | grad 0.06 | 1.61 it/s | ETA 8.61h
+step 10100/60000 | epoch 0 | loss 0.0138 | grad 0.09 | 1.61 it/s | ETA 8.59h
+step 10200/60000 | epoch 0 | loss 0.0151 | grad 0.03 | 1.62 it/s | ETA 8.56h
+step 10300/60000 | epoch 0 | loss 0.0153 | grad 0.05 | 1.62 it/s | ETA 8.54h
+step 10400/60000 | epoch 0 | loss 0.0136 | grad 0.09 | 1.62 it/s | ETA 8.52h
+step 10500/60000 | epoch 0 | loss 0.0151 | grad 0.12 | 1.62 it/s | ETA 8.49h
+step 10600/60000 | epoch 0 | loss 0.0150 | grad 0.10 | 1.62 it/s | ETA 8.47h
+step 10700/60000 | epoch 0 | loss 0.0151 | grad 0.10 | 1.62 it/s | ETA 8.45h
+step 10800/60000 | epoch 0 | loss 0.0159 | grad 0.12 | 1.62 it/s | ETA 8.42h
+step 10900/60000 | epoch 0 | loss 0.0157 | grad 0.09 | 1.62 it/s | ETA 8.40h
+step 11000/60000 | epoch 0 | loss 0.0146 | grad 0.08 | 1.62 it/s | ETA 8.38h
+step 11100/60000 | epoch 0 | loss 0.0146 | grad 0.08 | 1.63 it/s | ETA 8.35h
+step 11200/60000 | epoch 0 | loss 0.0171 | grad 0.10 | 1.63 it/s | ETA 8.33h
+step 11300/60000 | epoch 0 | loss 0.0154 | grad 0.09 | 1.63 it/s | ETA 8.31h
+step 11400/60000 | epoch 0 | loss 0.0151 | grad 0.05 | 1.63 it/s | ETA 8.29h
+step 11500/60000 | epoch 0 | loss 0.0150 | grad 0.06 | 1.63 it/s | ETA 8.26h
+step 11600/60000 | epoch 0 | loss 0.0152 | grad 0.13 | 1.63 it/s | ETA 8.24h
+step 11700/60000 | epoch 0 | loss 0.0154 | grad 0.10 | 1.63 it/s | ETA 8.22h
+step 11800/60000 | epoch 0 | loss 0.0150 | grad 0.05 | 1.63 it/s | ETA 8.20h
+step 11900/60000 | epoch 0 | loss 0.0149 | grad 0.20 | 1.63 it/s | ETA 8.17h
+step 12000/60000 | epoch 0 | loss 0.0166 | grad 0.02 | 1.64 it/s | ETA 8.15h
+step 12100/60000 | epoch 1 | loss 0.0149 | grad 0.07 | 1.63 it/s | ETA 8.14h
+step 12200/60000 | epoch 1 | loss 0.0145 | grad 0.08 | 1.63 it/s | ETA 8.12h
+step 12300/60000 | epoch 1 | loss 0.0139 | grad 0.05 | 1.64 it/s | ETA 8.10h
+step 12400/60000 | epoch 1 | loss 0.0151 | grad 0.20 | 1.64 it/s | ETA 8.08h
+step 12500/60000 | epoch 1 | loss 0.0156 | grad 0.08 | 1.64 it/s | ETA 8.06h
+step 12600/60000 | epoch 1 | loss 0.0169 | grad 0.08 | 1.64 it/s | ETA 8.03h
+step 12700/60000 | epoch 1 | loss 0.0148 | grad 0.08 | 1.64 it/s | ETA 8.01h
+step 12800/60000 | epoch 1 | loss 0.0150 | grad 0.12 | 1.64 it/s | ETA 7.99h
+step 12900/60000 | epoch 1 | loss 0.0157 | grad 0.05 | 1.64 it/s | ETA 7.97h
+step 13000/60000 | epoch 1 | loss 0.0144 | grad 0.08 | 1.64 it/s | ETA 7.95h
+step 13100/60000 | epoch 1 | loss 0.0160 | grad 0.08 | 1.64 it/s | ETA 7.93h
+step 13200/60000 | epoch 1 | loss 0.0151 | grad 0.03 | 1.64 it/s | ETA 7.91h
+step 13300/60000 | epoch 1 | loss 0.0155 | grad 0.07 | 1.64 it/s | ETA 7.89h
+step 13400/60000 | epoch 1 | loss 0.0142 | grad 0.11 | 1.65 it/s | ETA 7.87h
+step 13500/60000 | epoch 1 | loss 0.0158 | grad 0.05 | 1.65 it/s | ETA 7.85h
+step 13600/60000 | epoch 1 | loss 0.0145 | grad 0.09 | 1.65 it/s | ETA 7.83h
+step 13700/60000 | epoch 1 | loss 0.0139 | grad 0.06 | 1.65 it/s | ETA 7.80h
+step 13800/60000 | epoch 1 | loss 0.0143 | grad 0.11 | 1.65 it/s | ETA 7.78h
+step 13900/60000 | epoch 1 | loss 0.0140 | grad 0.13 | 1.65 it/s | ETA 7.76h
+step 14000/60000 | epoch 1 | loss 0.0136 | grad 0.06 | 1.65 it/s | ETA 7.74h
+step 14100/60000 | epoch 1 | loss 0.0151 | grad 0.14 | 1.65 it/s | ETA 7.72h
+step 14200/60000 | epoch 1 | loss 0.0127 | grad 0.09 | 1.65 it/s | ETA 7.70h
+step 14300/60000 | epoch 1 | loss 0.0151 | grad 0.09 | 1.65 it/s | ETA 7.68h
+step 14400/60000 | epoch 1 | loss 0.0150 | grad 0.08 | 1.65 it/s | ETA 7.66h
+step 14500/60000 | epoch 1 | loss 0.0144 | grad 0.10 | 1.65 it/s | ETA 7.64h
+step 14600/60000 | epoch 1 | loss 0.0146 | grad 0.03 | 1.65 it/s | ETA 7.62h
+step 14700/60000 | epoch 1 | loss 0.0155 | grad 0.11 | 1.66 it/s | ETA 7.60h
+step 14800/60000 | epoch 1 | loss 0.0157 | grad 0.08 | 1.66 it/s | ETA 7.58h
+step 14900/60000 | epoch 1 | loss 0.0139 | grad 0.09 | 1.66 it/s | ETA 7.56h
+step 15000/60000 | epoch 1 | loss 0.0139 | grad 0.07 | 1.66 it/s | ETA 7.55h
+step 15100/60000 | epoch 1 | loss 0.0146 | grad 0.05 | 1.65 it/s | ETA 7.55h
+step 15200/60000 | epoch 1 | loss 0.0140 | grad 0.10 | 1.65 it/s | ETA 7.53h
+step 15300/60000 | epoch 1 | loss 0.0143 | grad 0.06 | 1.65 it/s | ETA 7.52h
+step 15400/60000 | epoch 1 | loss 0.0143 | grad 0.04 | 1.65 it/s | ETA 7.51h
+step 15500/60000 | epoch 1 | loss 0.0137 | grad 0.13 | 1.65 it/s | ETA 7.50h
+step 15600/60000 | epoch 1 | loss 0.0141 | grad 0.12 | 1.65 it/s | ETA 7.48h
+step 15700/60000 | epoch 1 | loss 0.0138 | grad 0.11 | 1.65 it/s | ETA 7.47h
+step 15800/60000 | epoch 1 | loss 0.0148 | grad 0.03 | 1.65 it/s | ETA 7.46h
+step 15900/60000 | epoch 1 | loss 0.0147 | grad 0.09 | 1.65 it/s | ETA 7.44h
+step 16000/60000 | epoch 1 | loss 0.0157 | grad 0.07 | 1.65 it/s | ETA 7.43h
+step 16100/60000 | epoch 1 | loss 0.0143 | grad 0.06 | 1.64 it/s | ETA 7.41h
+step 16200/60000 | epoch 1 | loss 0.0166 | grad 0.05 | 1.64 it/s | ETA 7.40h
+step 16300/60000 | epoch 1 | loss 0.0139 | grad 0.07 | 1.64 it/s | ETA 7.38h
+step 16400/60000 | epoch 1 | loss 0.0149 | grad 0.10 | 1.64 it/s | ETA 7.37h
+step 16500/60000 | epoch 1 | loss 0.0141 | grad 0.04 | 1.64 it/s | ETA 7.35h
+step 16600/60000 | epoch 1 | loss 0.0149 | grad 0.11 | 1.64 it/s | ETA 7.34h
+step 16700/60000 | epoch 1 | loss 0.0141 | grad 0.10 | 1.64 it/s | ETA 7.32h
+step 16800/60000 | epoch 1 | loss 0.0143 | grad 0.08 | 1.64 it/s | ETA 7.31h
+step 16900/60000 | epoch 1 | loss 0.0136 | grad 0.12 | 1.64 it/s | ETA 7.29h
+step 17000/60000 | epoch 1 | loss 0.0150 | grad 0.03 | 1.64 it/s | ETA 7.28h
+step 17100/60000 | epoch 1 | loss 0.0148 | grad 0.14 | 1.64 it/s | ETA 7.26h
+step 17200/60000 | epoch 1 | loss 0.0153 | grad 0.04 | 1.64 it/s | ETA 7.25h
+step 17300/60000 | epoch 1 | loss 0.0143 | grad 0.11 | 1.64 it/s | ETA 7.23h
+step 17400/60000 | epoch 1 | loss 0.0150 | grad 0.08 | 1.64 it/s | ETA 7.22h
+step 17500/60000 | epoch 1 | loss 0.0142 | grad 0.12 | 1.64 it/s | ETA 7.20h
+step 17600/60000 | epoch 1 | loss 0.0139 | grad 0.04 | 1.64 it/s | ETA 7.18h
+step 17700/60000 | epoch 1 | loss 0.0140 | grad 0.06 | 1.64 it/s | ETA 7.17h
+step 17800/60000 | epoch 1 | loss 0.0125 | grad 0.07 | 1.64 it/s | ETA 7.15h
+step 17900/60000 | epoch 1 | loss 0.0163 | grad 0.09 | 1.64 it/s | ETA 7.14h
+step 18000/60000 | epoch 1 | loss 0.0129 | grad 0.03 | 1.64 it/s | ETA 7.12h
+step 18100/60000 | epoch 1 | loss 0.0153 | grad 0.06 | 1.64 it/s | ETA 7.12h
+step 18200/60000 | epoch 1 | loss 0.0165 | grad 0.07 | 1.63 it/s | ETA 7.10h
+step 18300/60000 | epoch 1 | loss 0.0141 | grad 0.09 | 1.63 it/s | ETA 7.09h
+step 18400/60000 | epoch 1 | loss 0.0150 | grad 0.08 | 1.63 it/s | ETA 7.07h
+step 18500/60000 | epoch 1 | loss 0.0159 | grad 0.12 | 1.63 it/s | ETA 7.06h
+step 18600/60000 | epoch 1 | loss 0.0125 | grad 0.10 | 1.63 it/s | ETA 7.04h
+step 18700/60000 | epoch 1 | loss 0.0140 | grad 0.04 | 1.63 it/s | ETA 7.03h
+step 18800/60000 | epoch 1 | loss 0.0130 | grad 0.08 | 1.63 it/s | ETA 7.01h
+step 18900/60000 | epoch 1 | loss 0.0153 | grad 0.09 | 1.63 it/s | ETA 6.99h
+step 19000/60000 | epoch 1 | loss 0.0163 | grad 0.08 | 1.63 it/s | ETA 6.98h
+step 19100/60000 | epoch 1 | loss 0.0131 | grad 0.04 | 1.63 it/s | ETA 6.96h
+step 19200/60000 | epoch 1 | loss 0.0148 | grad 0.12 | 1.63 it/s | ETA 6.95h
+step 19300/60000 | epoch 1 | loss 0.0130 | grad 0.09 | 1.63 it/s | ETA 6.93h
+step 19400/60000 | epoch 1 | loss 0.0140 | grad 0.09 | 1.63 it/s | ETA 6.91h
+step 19500/60000 | epoch 1 | loss 0.0138 | grad 0.05 | 1.63 it/s | ETA 6.90h
+step 19600/60000 | epoch 1 | loss 0.0150 | grad 0.06 | 1.63 it/s | ETA 6.88h
+step 19700/60000 | epoch 1 | loss 0.0146 | grad 0.08 | 1.63 it/s | ETA 6.87h
+step 19800/60000 | epoch 1 | loss 0.0133 | grad 0.07 | 1.63 it/s | ETA 6.85h
+step 19900/60000 | epoch 1 | loss 0.0137 | grad 0.08 | 1.63 it/s | ETA 6.83h
+step 20000/60000 | epoch 1 | loss 0.0146 | grad 0.09 | 1.63 it/s | ETA 6.80h
+step 20100/60000 | epoch 1 | loss 0.0144 | grad 0.11 | 1.63 it/s | ETA 6.78h
+step 20200/60000 | epoch 1 | loss 0.0151 | grad 0.03 | 1.63 it/s | ETA 6.76h
+step 20300/60000 | epoch 1 | loss 0.0148 | grad 0.06 | 1.64 it/s | ETA 6.74h
+step 20400/60000 | epoch 1 | loss 0.0144 | grad 0.07 | 1.64 it/s | ETA 6.72h
+step 20500/60000 | epoch 1 | loss 0.0141 | grad 0.09 | 1.64 it/s | ETA 6.70h
+step 20600/60000 | epoch 1 | loss 0.0142 | grad 0.10 | 1.64 it/s | ETA 6.68h
+step 20700/60000 | epoch 1 | loss 0.0132 | grad 0.08 | 1.64 it/s | ETA 6.66h
+step 20800/60000 | epoch 1 | loss 0.0137 | grad 0.06 | 1.64 it/s | ETA 6.64h
+step 20900/60000 | epoch 1 | loss 0.0131 | grad 0.02 | 1.64 it/s | ETA 6.62h
+step 21000/60000 | epoch 1 | loss 0.0132 | grad 0.12 | 1.64 it/s | ETA 6.60h
+step 21100/60000 | epoch 1 | loss 0.0144 | grad 0.09 | 1.64 it/s | ETA 6.59h
+step 21200/60000 | epoch 1 | loss 0.0133 | grad 0.04 | 1.64 it/s | ETA 6.57h
+step 21300/60000 | epoch 1 | loss 0.0127 | grad 0.11 | 1.64 it/s | ETA 6.55h
+step 21400/60000 | epoch 1 | loss 0.0139 | grad 0.02 | 1.64 it/s | ETA 6.53h
+step 21500/60000 | epoch 1 | loss 0.0145 | grad 0.10 | 1.64 it/s | ETA 6.51h
+step 21600/60000 | epoch 1 | loss 0.0128 | grad 0.03 | 1.64 it/s | ETA 6.49h
+step 21700/60000 | epoch 1 | loss 0.0143 | grad 0.06 | 1.64 it/s | ETA 6.47h
+step 21800/60000 | epoch 1 | loss 0.0153 | grad 0.07 | 1.65 it/s | ETA 6.45h
+step 21900/60000 | epoch 1 | loss 0.0140 | grad 0.03 | 1.65 it/s | ETA 6.43h
+step 22000/60000 | epoch 1 | loss 0.0152 | grad 0.04 | 1.65 it/s | ETA 6.41h
+step 22100/60000 | epoch 1 | loss 0.0137 | grad 0.06 | 1.65 it/s | ETA 6.39h
+step 22200/60000 | epoch 1 | loss 0.0137 | grad 0.08 | 1.65 it/s | ETA 6.37h
+step 22300/60000 | epoch 1 | loss 0.0142 | grad 0.07 | 1.65 it/s | ETA 6.35h
+step 22400/60000 | epoch 1 | loss 0.0133 | grad 0.11 | 1.65 it/s | ETA 6.33h
+step 22500/60000 | epoch 1 | loss 0.0128 | grad 0.10 | 1.65 it/s | ETA 6.31h
+step 22600/60000 | epoch 1 | loss 0.0131 | grad 0.09 | 1.65 it/s | ETA 6.29h
+step 22700/60000 | epoch 1 | loss 0.0155 | grad 0.07 | 1.65 it/s | ETA 6.27h
+step 22800/60000 | epoch 1 | loss 0.0140 | grad 0.13 | 1.65 it/s | ETA 6.25h
+step 22900/60000 | epoch 1 | loss 0.0136 | grad 0.09 | 1.65 it/s | ETA 6.23h
+step 23000/60000 | epoch 1 | loss 0.0151 | grad 0.10 | 1.65 it/s | ETA 6.21h
+step 23100/60000 | epoch 1 | loss 0.0146 | grad 0.10 | 1.65 it/s | ETA 6.19h
+step 23200/60000 | epoch 1 | loss 0.0125 | grad 0.06 | 1.66 it/s | ETA 6.18h
+step 23300/60000 | epoch 1 | loss 0.0132 | grad 0.07 | 1.66 it/s | ETA 6.16h
+step 23400/60000 | epoch 1 | loss 0.0144 | grad 0.18 | 1.66 it/s | ETA 6.14h
+step 23500/60000 | epoch 1 | loss 0.0134 | grad 0.10 | 1.66 it/s | ETA 6.12h
+step 23600/60000 | epoch 1 | loss 0.0148 | grad 0.08 | 1.66 it/s | ETA 6.10h
+step 23700/60000 | epoch 1 | loss 0.0147 | grad 0.07 | 1.66 it/s | ETA 6.08h
+step 23800/60000 | epoch 1 | loss 0.0138 | grad 0.02 | 1.66 it/s | ETA 6.06h
+step 23900/60000 | epoch 1 | loss 0.0138 | grad 0.06 | 1.66 it/s | ETA 6.04h
+step 24000/60000 | epoch 1 | loss 0.0132 | grad 0.10 | 1.66 it/s | ETA 6.02h
+step 24100/60000 | epoch 2 | loss 0.0144 | grad 0.08 | 1.66 it/s | ETA 6.01h
+step 24200/60000 | epoch 2 | loss 0.0149 | grad 0.09 | 1.66 it/s | ETA 5.99h
+step 24300/60000 | epoch 2 | loss 0.0148 | grad 0.04 | 1.66 it/s | ETA 5.97h
+step 24400/60000 | epoch 2 | loss 0.0144 | grad 0.08 | 1.66 it/s | ETA 5.95h
+step 24500/60000 | epoch 2 | loss 0.0140 | grad 0.03 | 1.66 it/s | ETA 5.93h
+step 24600/60000 | epoch 2 | loss 0.0139 | grad 0.02 | 1.66 it/s | ETA 5.91h
+step 24700/60000 | epoch 2 | loss 0.0143 | grad 0.08 | 1.66 it/s | ETA 5.89h
+step 24800/60000 | epoch 2 | loss 0.0137 | grad 0.05 | 1.66 it/s | ETA 5.88h
+step 24900/60000 | epoch 2 | loss 0.0152 | grad 0.09 | 1.66 it/s | ETA 5.86h
+step 25000/60000 | epoch 2 | loss 0.0149 | grad 0.06 | 1.67 it/s | ETA 5.84h
+step 25100/60000 | epoch 2 | loss 0.0141 | grad 0.07 | 1.67 it/s | ETA 5.82h
+step 25200/60000 | epoch 2 | loss 0.0145 | grad 0.06 | 1.67 it/s | ETA 5.80h
+step 25300/60000 | epoch 2 | loss 0.0137 | grad 0.07 | 1.67 it/s | ETA 5.78h
+step 25400/60000 | epoch 2 | loss 0.0149 | grad 0.10 | 1.67 it/s | ETA 5.76h
+step 25500/60000 | epoch 2 | loss 0.0144 | grad 0.09 | 1.67 it/s | ETA 5.75h
+step 25600/60000 | epoch 2 | loss 0.0146 | grad 0.05 | 1.67 it/s | ETA 5.73h
+step 25700/60000 | epoch 2 | loss 0.0150 | grad 0.12 | 1.67 it/s | ETA 5.71h
+step 25800/60000 | epoch 2 | loss 0.0133 | grad 0.09 | 1.67 it/s | ETA 5.69h
+step 25900/60000 | epoch 2 | loss 0.0132 | grad 0.13 | 1.67 it/s | ETA 5.67h
+step 26000/60000 | epoch 2 | loss 0.0124 | grad 0.11 | 1.67 it/s | ETA 5.65h
+step 26100/60000 | epoch 2 | loss 0.0152 | grad 0.10 | 1.67 it/s | ETA 5.63h
+step 26200/60000 | epoch 2 | loss 0.0147 | grad 0.08 | 1.67 it/s | ETA 5.62h
+step 26300/60000 | epoch 2 | loss 0.0155 | grad 0.06 | 1.67 it/s | ETA 5.60h
+step 26400/60000 | epoch 2 | loss 0.0119 | grad 0.10 | 1.67 it/s | ETA 5.58h
+step 26500/60000 | epoch 2 | loss 0.0132 | grad 0.11 | 1.67 it/s | ETA 5.56h
+step 26600/60000 | epoch 2 | loss 0.0135 | grad 0.06 | 1.67 it/s | ETA 5.54h
+step 26700/60000 | epoch 2 | loss 0.0136 | grad 0.05 | 1.67 it/s | ETA 5.52h
+step 26800/60000 | epoch 2 | loss 0.0136 | grad 0.05 | 1.67 it/s | ETA 5.51h
+step 26900/60000 | epoch 2 | loss 0.0138 | grad 0.06 | 1.68 it/s | ETA 5.49h
+step 27000/60000 | epoch 2 | loss 0.0143 | grad 0.06 | 1.68 it/s | ETA 5.47h
+step 27100/60000 | epoch 2 | loss 0.0143 | grad 0.12 | 1.68 it/s | ETA 5.46h
+step 27200/60000 | epoch 2 | loss 0.0149 | grad 0.13 | 1.68 it/s | ETA 5.44h
+step 27300/60000 | epoch 2 | loss 0.0135 | grad 0.08 | 1.68 it/s | ETA 5.42h
+step 27400/60000 | epoch 2 | loss 0.0152 | grad 0.03 | 1.68 it/s | ETA 5.40h
+step 27500/60000 | epoch 2 | loss 0.0124 | grad 0.05 | 1.68 it/s | ETA 5.38h
+step 27600/60000 | epoch 2 | loss 0.0152 | grad 0.06 | 1.68 it/s | ETA 5.36h
+step 27700/60000 | epoch 2 | loss 0.0137 | grad 0.07 | 1.68 it/s | ETA 5.35h
+step 27800/60000 | epoch 2 | loss 0.0131 | grad 0.09 | 1.68 it/s | ETA 5.33h
+step 27900/60000 | epoch 2 | loss 0.0142 | grad 0.13 | 1.68 it/s | ETA 5.31h
+step 28000/60000 | epoch 2 | loss 0.0146 | grad 0.09 | 1.68 it/s | ETA 5.29h
+step 28100/60000 | epoch 2 | loss 0.0147 | grad 0.05 | 1.68 it/s | ETA 5.27h
+step 28200/60000 | epoch 2 | loss 0.0128 | grad 0.07 | 1.68 it/s | ETA 5.26h
+step 28300/60000 | epoch 2 | loss 0.0141 | grad 0.05 | 1.68 it/s | ETA 5.24h
+step 28400/60000 | epoch 2 | loss 0.0142 | grad 0.12 | 1.68 it/s | ETA 5.22h
+step 28500/60000 | epoch 2 | loss 0.0141 | grad 0.10 | 1.68 it/s | ETA 5.20h
+step 28600/60000 | epoch 2 | loss 0.0144 | grad 0.12 | 1.68 it/s | ETA 5.18h
+step 28700/60000 | epoch 2 | loss 0.0130 | grad 0.14 | 1.68 it/s | ETA 5.17h
+step 28800/60000 | epoch 2 | loss 0.0151 | grad 0.03 | 1.68 it/s | ETA 5.15h
+step 28900/60000 | epoch 2 | loss 0.0143 | grad 0.06 | 1.68 it/s | ETA 5.13h
+step 29000/60000 | epoch 2 | loss 0.0147 | grad 0.05 | 1.68 it/s | ETA 5.11h
+step 29100/60000 | epoch 2 | loss 0.0139 | grad 0.12 | 1.68 it/s | ETA 5.10h
+step 29200/60000 | epoch 2 | loss 0.0137 | grad 0.09 | 1.68 it/s | ETA 5.08h
+step 29300/60000 | epoch 2 | loss 0.0133 | grad 0.09 | 1.69 it/s | ETA 5.06h
+step 29400/60000 | epoch 2 | loss 0.0140 | grad 0.04 | 1.69 it/s | ETA 5.04h
+step 29500/60000 | epoch 2 | loss 0.0134 | grad 0.06 | 1.69 it/s | ETA 5.02h
+step 29600/60000 | epoch 2 | loss 0.0133 | grad 0.04 | 1.69 it/s | ETA 5.01h
+step 29700/60000 | epoch 2 | loss 0.0131 | grad 0.05 | 1.69 it/s | ETA 4.99h
+step 29800/60000 | epoch 2 | loss 0.0140 | grad 0.08 | 1.69 it/s | ETA 4.97h
+step 29900/60000 | epoch 2 | loss 0.0154 | grad 0.08 | 1.69 it/s | ETA 4.95h
+step 30000/60000 | epoch 2 | loss 0.0144 | grad 0.10 | 1.69 it/s | ETA 4.94h
+step 30100/60000 | epoch 2 | loss 0.0142 | grad 0.11 | 1.69 it/s | ETA 4.92h
+step 30200/60000 | epoch 2 | loss 0.0137 | grad 0.07 | 1.69 it/s | ETA 4.90h
+step 30300/60000 | epoch 2 | loss 0.0134 | grad 0.04 | 1.69 it/s | ETA 4.89h
+step 30400/60000 | epoch 2 | loss 0.0144 | grad 0.06 | 1.69 it/s | ETA 4.87h
+step 30500/60000 | epoch 2 | loss 0.0152 | grad 0.03 | 1.69 it/s | ETA 4.85h
+step 30600/60000 | epoch 2 | loss 0.0147 | grad 0.09 | 1.69 it/s | ETA 4.83h
+step 30700/60000 | epoch 2 | loss 0.0151 | grad 0.03 | 1.69 it/s | ETA 4.82h
+step 30800/60000 | epoch 2 | loss 0.0145 | grad 0.04 | 1.69 it/s | ETA 4.80h
+step 30900/60000 | epoch 2 | loss 0.0150 | grad 0.11 | 1.69 it/s | ETA 4.78h
+step 31000/60000 | epoch 2 | loss 0.0137 | grad 0.10 | 1.69 it/s | ETA 4.76h
+step 31100/60000 | epoch 2 | loss 0.0132 | grad 0.04 | 1.69 it/s | ETA 4.74h
+step 31200/60000 | epoch 2 | loss 0.0149 | grad 0.06 | 1.69 it/s | ETA 4.73h
+step 31300/60000 | epoch 2 | loss 0.0136 | grad 0.06 | 1.69 it/s | ETA 4.71h
+step 31400/60000 | epoch 2 | loss 0.0133 | grad 0.06 | 1.69 it/s | ETA 4.69h
+step 31500/60000 | epoch 2 | loss 0.0152 | grad 0.15 | 1.69 it/s | ETA 4.68h
+step 31600/60000 | epoch 2 | loss 0.0142 | grad 0.08 | 1.69 it/s | ETA 4.66h
+step 31700/60000 | epoch 2 | loss 0.0132 | grad 0.06 | 1.69 it/s | ETA 4.64h
+step 31800/60000 | epoch 2 | loss 0.0142 | grad 0.10 | 1.69 it/s | ETA 4.62h
+step 31900/60000 | epoch 2 | loss 0.0150 | grad 0.08 | 1.69 it/s | ETA 4.61h
+step 32000/60000 | epoch 2 | loss 0.0130 | grad 0.02 | 1.70 it/s | ETA 4.59h
+step 32100/60000 | epoch 2 | loss 0.0149 | grad 0.05 | 1.70 it/s | ETA 4.57h
+step 32200/60000 | epoch 2 | loss 0.0125 | grad 0.10 | 1.70 it/s | ETA 4.55h
+step 32300/60000 | epoch 2 | loss 0.0138 | grad 0.07 | 1.70 it/s | ETA 4.54h
+step 32400/60000 | epoch 2 | loss 0.0140 | grad 0.05 | 1.70 it/s | ETA 4.52h
+step 32500/60000 | epoch 2 | loss 0.0168 | grad 0.06 | 1.70 it/s | ETA 4.50h
+step 32600/60000 | epoch 2 | loss 0.0147 | grad 0.08 | 1.70 it/s | ETA 4.48h
+step 32700/60000 | epoch 2 | loss 0.0143 | grad 0.06 | 1.70 it/s | ETA 4.47h
+step 32800/60000 | epoch 2 | loss 0.0143 | grad 0.05 | 1.70 it/s | ETA 4.45h
+step 32900/60000 | epoch 2 | loss 0.0132 | grad 0.06 | 1.70 it/s | ETA 4.43h
+step 33000/60000 | epoch 2 | loss 0.0140 | grad 0.03 | 1.70 it/s | ETA 4.41h
+step 33100/60000 | epoch 2 | loss 0.0143 | grad 0.13 | 1.70 it/s | ETA 4.40h
+step 33200/60000 | epoch 2 | loss 0.0155 | grad 0.04 | 1.70 it/s | ETA 4.38h
+step 33300/60000 | epoch 2 | loss 0.0147 | grad 0.07 | 1.70 it/s | ETA 4.37h
+step 33400/60000 | epoch 2 | loss 0.0128 | grad 0.08 | 1.70 it/s | ETA 4.35h
+step 33500/60000 | epoch 2 | loss 0.0125 | grad 0.12 | 1.70 it/s | ETA 4.33h
+step 33600/60000 | epoch 2 | loss 0.0131 | grad 0.07 | 1.70 it/s | ETA 4.31h
+step 33700/60000 | epoch 2 | loss 0.0142 | grad 0.05 | 1.70 it/s | ETA 4.30h
+step 33800/60000 | epoch 2 | loss 0.0136 | grad 0.08 | 1.70 it/s | ETA 4.28h
+step 33900/60000 | epoch 2 | loss 0.0131 | grad 0.11 | 1.70 it/s | ETA 4.26h
+step 34000/60000 | epoch 2 | loss 0.0135 | grad 0.05 | 1.70 it/s | ETA 4.24h
+step 34100/60000 | epoch 2 | loss 0.0137 | grad 0.09 | 1.70 it/s | ETA 4.23h
+step 34200/60000 | epoch 2 | loss 0.0120 | grad 0.05 | 1.70 it/s | ETA 4.21h
+step 34300/60000 | epoch 2 | loss 0.0142 | grad 0.04 | 1.70 it/s | ETA 4.19h
+step 34400/60000 | epoch 2 | loss 0.0137 | grad 0.06 | 1.70 it/s | ETA 4.18h
+step 34500/60000 | epoch 2 | loss 0.0137 | grad 0.11 | 1.70 it/s | ETA 4.16h
+step 34600/60000 | epoch 2 | loss 0.0141 | grad 0.02 | 1.70 it/s | ETA 4.14h
+step 34700/60000 | epoch 2 | loss 0.0145 | grad 0.10 | 1.70 it/s | ETA 4.13h
+step 34800/60000 | epoch 2 | loss 0.0152 | grad 0.08 | 1.70 it/s | ETA 4.11h
+step 34900/60000 | epoch 2 | loss 0.0133 | grad 0.05 | 1.70 it/s | ETA 4.09h
+step 35000/60000 | epoch 2 | loss 0.0137 | grad 0.09 | 1.70 it/s | ETA 4.07h
+step 35100/60000 | epoch 2 | loss 0.0143 | grad 0.09 | 1.70 it/s | ETA 4.06h
+step 35200/60000 | epoch 2 | loss 0.0136 | grad 0.05 | 1.71 it/s | ETA 4.04h
+step 35300/60000 | epoch 2 | loss 0.0129 | grad 0.06 | 1.71 it/s | ETA 4.02h
+step 35400/60000 | epoch 2 | loss 0.0141 | grad 0.11 | 1.71 it/s | ETA 4.01h
+step 35500/60000 | epoch 2 | loss 0.0127 | grad 0.08 | 1.71 it/s | ETA 3.99h
+step 35600/60000 | epoch 2 | loss 0.0132 | grad 0.05 | 1.71 it/s | ETA 3.97h
+step 35700/60000 | epoch 2 | loss 0.0148 | grad 0.04 | 1.71 it/s | ETA 3.95h
+step 35800/60000 | epoch 2 | loss 0.0132 | grad 0.08 | 1.71 it/s | ETA 3.94h
+step 35900/60000 | epoch 2 | loss 0.0144 | grad 0.08 | 1.71 it/s | ETA 3.92h
+step 36000/60000 | epoch 2 | loss 0.0127 | grad 0.03 | 1.71 it/s | ETA 3.90h
+step 36100/60000 | epoch 3 | loss 0.0140 | grad 0.08 | 1.71 it/s | ETA 3.89h
+step 36200/60000 | epoch 3 | loss 0.0142 | grad 0.04 | 1.71 it/s | ETA 3.87h
+step 36300/60000 | epoch 3 | loss 0.0125 | grad 0.07 | 1.71 it/s | ETA 3.86h
+step 36400/60000 | epoch 3 | loss 0.0134 | grad 0.06 | 1.71 it/s | ETA 3.84h
+step 36500/60000 | epoch 3 | loss 0.0138 | grad 0.06 | 1.71 it/s | ETA 3.82h
+step 36600/60000 | epoch 3 | loss 0.0144 | grad 0.05 | 1.71 it/s | ETA 3.80h
+step 36700/60000 | epoch 3 | loss 0.0143 | grad 0.03 | 1.71 it/s | ETA 3.79h
+step 36800/60000 | epoch 3 | loss 0.0146 | grad 0.09 | 1.71 it/s | ETA 3.77h
+step 36900/60000 | epoch 3 | loss 0.0148 | grad 0.04 | 1.71 it/s | ETA 3.75h
+step 37000/60000 | epoch 3 | loss 0.0138 | grad 0.02 | 1.71 it/s | ETA 3.74h
+step 37100/60000 | epoch 3 | loss 0.0133 | grad 0.06 | 1.71 it/s | ETA 3.72h
+step 37200/60000 | epoch 3 | loss 0.0138 | grad 0.04 | 1.71 it/s | ETA 3.70h
+step 37300/60000 | epoch 3 | loss 0.0133 | grad 0.10 | 1.71 it/s | ETA 3.69h
+step 37400/60000 | epoch 3 | loss 0.0145 | grad 0.08 | 1.71 it/s | ETA 3.67h
+step 37500/60000 | epoch 3 | loss 0.0135 | grad 0.05 | 1.71 it/s | ETA 3.65h
+step 37600/60000 | epoch 3 | loss 0.0127 | grad 0.07 | 1.71 it/s | ETA 3.64h
+step 37700/60000 | epoch 3 | loss 0.0138 | grad 0.10 | 1.71 it/s | ETA 3.62h
+step 37800/60000 | epoch 3 | loss 0.0145 | grad 0.09 | 1.71 it/s | ETA 3.60h
+step 37900/60000 | epoch 3 | loss 0.0134 | grad 0.05 | 1.71 it/s | ETA 3.59h
+step 38000/60000 | epoch 3 | loss 0.0130 | grad 0.11 | 1.71 it/s | ETA 3.57h
+step 38100/60000 | epoch 3 | loss 0.0128 | grad 0.06 | 1.71 it/s | ETA 3.55h
+step 38200/60000 | epoch 3 | loss 0.0137 | grad 0.04 | 1.71 it/s | ETA 3.54h
+step 38300/60000 | epoch 3 | loss 0.0137 | grad 0.08 | 1.71 it/s | ETA 3.52h
+step 38400/60000 | epoch 3 | loss 0.0140 | grad 0.06 | 1.71 it/s | ETA 3.50h
+step 38500/60000 | epoch 3 | loss 0.0130 | grad 0.05 | 1.71 it/s | ETA 3.48h
+step 38600/60000 | epoch 3 | loss 0.0145 | grad 0.06 | 1.71 it/s | ETA 3.47h
+step 38700/60000 | epoch 3 | loss 0.0148 | grad 0.05 | 1.71 it/s | ETA 3.45h
+step 38800/60000 | epoch 3 | loss 0.0140 | grad 0.05 | 1.71 it/s | ETA 3.43h
+step 38900/60000 | epoch 3 | loss 0.0133 | grad 0.11 | 1.71 it/s | ETA 3.42h
+step 39000/60000 | epoch 3 | loss 0.0133 | grad 0.08 | 1.72 it/s | ETA 3.40h
+step 39100/60000 | epoch 3 | loss 0.0143 | grad 0.02 | 1.71 it/s | ETA 3.39h
+step 39200/60000 | epoch 3 | loss 0.0149 | grad 0.08 | 1.71 it/s | ETA 3.37h
+step 39300/60000 | epoch 3 | loss 0.0131 | grad 0.11 | 1.72 it/s | ETA 3.35h
+step 39400/60000 | epoch 3 | loss 0.0148 | grad 0.03 | 1.72 it/s | ETA 3.34h
+step 39500/60000 | epoch 3 | loss 0.0125 | grad 0.07 | 1.72 it/s | ETA 3.32h
+step 39600/60000 | epoch 3 | loss 0.0146 | grad 0.09 | 1.72 it/s | ETA 3.30h
+step 39700/60000 | epoch 3 | loss 0.0137 | grad 0.11 | 1.72 it/s | ETA 3.29h
+step 39800/60000 | epoch 3 | loss 0.0154 | grad 0.12 | 1.72 it/s | ETA 3.27h
+step 39900/60000 | epoch 3 | loss 0.0140 | grad 0.06 | 1.72 it/s | ETA 3.25h
+step 40000/60000 | epoch 3 | loss 0.0141 | grad 0.07 | 1.72 it/s | ETA 3.24h
+step 40100/60000 | epoch 3 | loss 0.0133 | grad 0.11 | 1.72 it/s | ETA 3.22h
+step 40200/60000 | epoch 3 | loss 0.0132 | grad 0.05 | 1.72 it/s | ETA 3.20h
+step 40300/60000 | epoch 3 | loss 0.0131 | grad 0.08 | 1.72 it/s | ETA 3.19h
+step 40400/60000 | epoch 3 | loss 0.0140 | grad 0.06 | 1.72 it/s | ETA 3.17h
+step 40500/60000 | epoch 3 | loss 0.0150 | grad 0.08 | 1.72 it/s | ETA 3.15h
+step 40600/60000 | epoch 3 | loss 0.0127 | grad 0.08 | 1.72 it/s | ETA 3.14h
+step 40700/60000 | epoch 3 | loss 0.0122 | grad 0.06 | 1.72 it/s | ETA 3.12h
+step 40800/60000 | epoch 3 | loss 0.0135 | grad 0.09 | 1.72 it/s | ETA 3.10h
+step 40900/60000 | epoch 3 | loss 0.0151 | grad 0.06 | 1.72 it/s | ETA 3.09h
+step 41000/60000 | epoch 3 | loss 0.0140 | grad 0.08 | 1.72 it/s | ETA 3.07h
+step 41100/60000 | epoch 3 | loss 0.0131 | grad 0.12 | 1.72 it/s | ETA 3.05h
+step 41200/60000 | epoch 3 | loss 0.0129 | grad 0.12 | 1.72 it/s | ETA 3.04h
+step 41300/60000 | epoch 3 | loss 0.0118 | grad 0.07 | 1.72 it/s | ETA 3.02h
+step 41400/60000 | epoch 3 | loss 0.0127 | grad 0.10 | 1.72 it/s | ETA 3.00h
+step 41500/60000 | epoch 3 | loss 0.0139 | grad 0.03 | 1.72 it/s | ETA 2.99h
+step 41600/60000 | epoch 3 | loss 0.0149 | grad 0.07 | 1.72 it/s | ETA 2.97h
+step 41700/60000 | epoch 3 | loss 0.0141 | grad 0.06 | 1.72 it/s | ETA 2.95h
+step 41800/60000 | epoch 3 | loss 0.0143 | grad 0.06 | 1.72 it/s | ETA 2.94h
+step 41900/60000 | epoch 3 | loss 0.0139 | grad 0.04 | 1.72 it/s | ETA 2.92h
+step 42000/60000 | epoch 3 | loss 0.0148 | grad 0.04 | 1.72 it/s | ETA 2.90h
+step 42100/60000 | epoch 3 | loss 0.0122 | grad 0.05 | 1.72 it/s | ETA 2.89h
+step 42200/60000 | epoch 3 | loss 0.0122 | grad 0.05 | 1.72 it/s | ETA 2.87h
+step 42300/60000 | epoch 3 | loss 0.0130 | grad 0.04 | 1.72 it/s | ETA 2.86h
+step 42400/60000 | epoch 3 | loss 0.0140 | grad 0.08 | 1.72 it/s | ETA 2.84h
+step 42500/60000 | epoch 3 | loss 0.0139 | grad 0.09 | 1.72 it/s | ETA 2.82h
+step 42600/60000 | epoch 3 | loss 0.0129 | grad 0.03 | 1.72 it/s | ETA 2.81h
+step 42700/60000 | epoch 3 | loss 0.0123 | grad 0.06 | 1.72 it/s | ETA 2.79h
+step 42800/60000 | epoch 3 | loss 0.0139 | grad 0.08 | 1.72 it/s | ETA 2.77h
+step 42900/60000 | epoch 3 | loss 0.0141 | grad 0.10 | 1.72 it/s | ETA 2.76h
+step 43000/60000 | epoch 3 | loss 0.0138 | grad 0.05 | 1.72 it/s | ETA 2.74h
+step 43100/60000 | epoch 3 | loss 0.0144 | grad 0.07 | 1.72 it/s | ETA 2.72h
+step 43200/60000 | epoch 3 | loss 0.0136 | grad 0.06 | 1.72 it/s | ETA 2.71h
+step 43300/60000 | epoch 3 | loss 0.0140 | grad 0.08 | 1.72 it/s | ETA 2.69h
+step 43400/60000 | epoch 3 | loss 0.0140 | grad 0.05 | 1.72 it/s | ETA 2.67h
+step 43500/60000 | epoch 3 | loss 0.0129 | grad 0.06 | 1.72 it/s | ETA 2.66h
+step 43600/60000 | epoch 3 | loss 0.0147 | grad 0.06 | 1.72 it/s | ETA 2.64h
+step 43700/60000 | epoch 3 | loss 0.0132 | grad 0.05 | 1.72 it/s | ETA 2.63h
+step 43800/60000 | epoch 3 | loss 0.0138 | grad 0.08 | 1.72 it/s | ETA 2.61h
+step 43900/60000 | epoch 3 | loss 0.0148 | grad 0.11 | 1.72 it/s | ETA 2.59h
+step 44000/60000 | epoch 3 | loss 0.0132 | grad 0.05 | 1.73 it/s | ETA 2.58h
+step 44100/60000 | epoch 3 | loss 0.0140 | grad 0.09 | 1.73 it/s | ETA 2.56h
+step 44200/60000 | epoch 3 | loss 0.0145 | grad 0.10 | 1.73 it/s | ETA 2.54h
+step 44300/60000 | epoch 3 | loss 0.0141 | grad 0.05 | 1.73 it/s | ETA 2.53h
+step 44400/60000 | epoch 3 | loss 0.0133 | grad 0.05 | 1.73 it/s | ETA 2.51h
+step 44500/60000 | epoch 3 | loss 0.0128 | grad 0.04 | 1.73 it/s | ETA 2.49h
+step 44600/60000 | epoch 3 | loss 0.0122 | grad 0.08 | 1.73 it/s | ETA 2.48h
+step 44700/60000 | epoch 3 | loss 0.0141 | grad 0.06 | 1.73 it/s | ETA 2.46h
+step 44800/60000 | epoch 3 | loss 0.0131 | grad 0.06 | 1.73 it/s | ETA 2.45h
+step 44900/60000 | epoch 3 | loss 0.0135 | grad 0.11 | 1.73 it/s | ETA 2.43h
+step 45000/60000 | epoch 3 | loss 0.0127 | grad 0.07 | 1.73 it/s | ETA 2.41h
+step 45100/60000 | epoch 3 | loss 0.0141 | grad 0.07 | 1.73 it/s | ETA 2.40h
+step 45200/60000 | epoch 3 | loss 0.0148 | grad 0.12 | 1.73 it/s | ETA 2.38h
+step 45300/60000 | epoch 3 | loss 0.0129 | grad 0.08 | 1.73 it/s | ETA 2.36h
+step 45400/60000 | epoch 3 | loss 0.0133 | grad 0.11 | 1.73 it/s | ETA 2.35h
+step 45500/60000 | epoch 3 | loss 0.0132 | grad 0.06 | 1.73 it/s | ETA 2.33h
+step 45600/60000 | epoch 3 | loss 0.0134 | grad 0.04 | 1.73 it/s | ETA 2.32h
+step 45700/60000 | epoch 3 | loss 0.0125 | grad 0.05 | 1.73 it/s | ETA 2.30h
+step 45800/60000 | epoch 3 | loss 0.0133 | grad 0.08 | 1.73 it/s | ETA 2.28h
+step 45900/60000 | epoch 3 | loss 0.0125 | grad 0.08 | 1.73 it/s | ETA 2.27h
+step 46000/60000 | epoch 3 | loss 0.0124 | grad 0.05 | 1.73 it/s | ETA 2.25h
+step 46100/60000 | epoch 3 | loss 0.0138 | grad 0.08 | 1.73 it/s | ETA 2.23h
+step 46200/60000 | epoch 3 | loss 0.0137 | grad 0.07 | 1.73 it/s | ETA 2.22h
+step 46300/60000 | epoch 3 | loss 0.0130 | grad 0.08 | 1.73 it/s | ETA 2.20h
+step 46400/60000 | epoch 3 | loss 0.0131 | grad 0.08 | 1.73 it/s | ETA 2.18h
+step 46500/60000 | epoch 3 | loss 0.0122 | grad 0.04 | 1.73 it/s | ETA 2.17h
+step 46600/60000 | epoch 3 | loss 0.0143 | grad 0.09 | 1.73 it/s | ETA 2.15h
+step 46700/60000 | epoch 3 | loss 0.0134 | grad 0.05 | 1.73 it/s | ETA 2.14h
+step 46800/60000 | epoch 3 | loss 0.0149 | grad 0.05 | 1.73 it/s | ETA 2.12h
+step 46900/60000 | epoch 3 | loss 0.0139 | grad 0.11 | 1.73 it/s | ETA 2.10h
+step 47000/60000 | epoch 3 | loss 0.0134 | grad 0.05 | 1.73 it/s | ETA 2.09h
+step 47100/60000 | epoch 3 | loss 0.0129 | grad 0.02 | 1.73 it/s | ETA 2.07h
+step 47200/60000 | epoch 3 | loss 0.0132 | grad 0.09 | 1.73 it/s | ETA 2.05h
+step 47300/60000 | epoch 3 | loss 0.0140 | grad 0.06 | 1.73 it/s | ETA 2.04h
+step 47400/60000 | epoch 3 | loss 0.0138 | grad 0.04 | 1.73 it/s | ETA 2.02h
+step 47500/60000 | epoch 3 | loss 0.0137 | grad 0.09 | 1.73 it/s | ETA 2.01h
+step 47600/60000 | epoch 3 | loss 0.0149 | grad 0.02 | 1.73 it/s | ETA 1.99h
+step 47700/60000 | epoch 3 | loss 0.0131 | grad 0.06 | 1.73 it/s | ETA 1.97h
+step 47800/60000 | epoch 3 | loss 0.0134 | grad 0.08 | 1.73 it/s | ETA 1.96h
+step 47900/60000 | epoch 3 | loss 0.0135 | grad 0.02 | 1.73 it/s | ETA 1.94h
+step 48000/60000 | epoch 3 | loss 0.0126 | grad 0.05 | 1.73 it/s | ETA 1.92h
+step 48100/60000 | epoch 3 | loss 0.0136 | grad 0.08 | 1.73 it/s | ETA 1.91h
+step 48200/60000 | epoch 4 | loss 0.0137 | grad 0.05 | 1.73 it/s | ETA 1.89h
+step 48300/60000 | epoch 4 | loss 0.0148 | grad 0.10 | 1.73 it/s | ETA 1.88h
+step 48400/60000 | epoch 4 | loss 0.0141 | grad 0.09 | 1.73 it/s | ETA 1.86h
+step 48500/60000 | epoch 4 | loss 0.0129 | grad 0.04 | 1.73 it/s | ETA 1.84h
+step 48600/60000 | epoch 4 | loss 0.0125 | grad 0.04 | 1.73 it/s | ETA 1.83h
+step 48700/60000 | epoch 4 | loss 0.0140 | grad 0.06 | 1.73 it/s | ETA 1.81h
+step 48800/60000 | epoch 4 | loss 0.0145 | grad 0.04 | 1.73 it/s | ETA 1.80h
+step 48900/60000 | epoch 4 | loss 0.0122 | grad 0.03 | 1.73 it/s | ETA 1.78h
+step 49000/60000 | epoch 4 | loss 0.0131 | grad 0.08 | 1.73 it/s | ETA 1.76h
+step 49100/60000 | epoch 4 | loss 0.0147 | grad 0.04 | 1.73 it/s | ETA 1.75h
+step 49200/60000 | epoch 4 | loss 0.0133 | grad 0.09 | 1.73 it/s | ETA 1.73h
+step 49300/60000 | epoch 4 | loss 0.0132 | grad 0.05 | 1.73 it/s | ETA 1.71h
+step 49400/60000 | epoch 4 | loss 0.0137 | grad 0.04 | 1.73 it/s | ETA 1.70h
+step 49500/60000 | epoch 4 | loss 0.0123 | grad 0.05 | 1.73 it/s | ETA 1.68h
+step 49600/60000 | epoch 4 | loss 0.0132 | grad 0.08 | 1.73 it/s | ETA 1.67h
+step 49700/60000 | epoch 4 | loss 0.0135 | grad 0.05 | 1.73 it/s | ETA 1.65h
+step 49800/60000 | epoch 4 | loss 0.0143 | grad 0.03 | 1.73 it/s | ETA 1.63h
+step 49900/60000 | epoch 4 | loss 0.0116 | grad 0.04 | 1.73 it/s | ETA 1.62h
+step 50000/60000 | epoch 4 | loss 0.0138 | grad 0.09 | 1.74 it/s | ETA 1.60h
+step 50100/60000 | epoch 4 | loss 0.0151 | grad 0.06 | 1.74 it/s | ETA 1.58h
+step 50200/60000 | epoch 4 | loss 0.0153 | grad 0.04 | 1.74 it/s | ETA 1.57h
+step 50300/60000 | epoch 4 | loss 0.0114 | grad 0.03 | 1.74 it/s | ETA 1.55h
+step 50400/60000 | epoch 4 | loss 0.0128 | grad 0.08 | 1.74 it/s | ETA 1.54h
+step 50500/60000 | epoch 4 | loss 0.0149 | grad 0.09 | 1.74 it/s | ETA 1.52h
+step 50600/60000 | epoch 4 | loss 0.0141 | grad 0.07 | 1.74 it/s | ETA 1.50h
+step 50700/60000 | epoch 4 | loss 0.0142 | grad 0.03 | 1.74 it/s | ETA 1.49h
+step 50800/60000 | epoch 4 | loss 0.0141 | grad 0.05 | 1.74 it/s | ETA 1.47h
+step 50900/60000 | epoch 4 | loss 0.0128 | grad 0.07 | 1.74 it/s | ETA 1.46h
+step 51000/60000 | epoch 4 | loss 0.0136 | grad 0.08 | 1.74 it/s | ETA 1.44h
+step 51100/60000 | epoch 4 | loss 0.0145 | grad 0.03 | 1.74 it/s | ETA 1.42h
+step 51200/60000 | epoch 4 | loss 0.0128 | grad 0.08 | 1.74 it/s | ETA 1.41h
+step 51300/60000 | epoch 4 | loss 0.0139 | grad 0.03 | 1.74 it/s | ETA 1.39h
+step 51400/60000 | epoch 4 | loss 0.0131 | grad 0.09 | 1.74 it/s | ETA 1.38h
+step 51500/60000 | epoch 4 | loss 0.0138 | grad 0.03 | 1.74 it/s | ETA 1.36h
+step 51600/60000 | epoch 4 | loss 0.0125 | grad 0.04 | 1.74 it/s | ETA 1.34h
+step 51700/60000 | epoch 4 | loss 0.0115 | grad 0.06 | 1.74 it/s | ETA 1.33h
+step 51800/60000 | epoch 4 | loss 0.0118 | grad 0.05 | 1.74 it/s | ETA 1.31h
+step 51900/60000 | epoch 4 | loss 0.0132 | grad 0.05 | 1.74 it/s | ETA 1.29h
+step 52000/60000 | epoch 4 | loss 0.0128 | grad 0.09 | 1.74 it/s | ETA 1.28h
+step 52100/60000 | epoch 4 | loss 0.0144 | grad 0.09 | 1.74 it/s | ETA 1.26h
+step 52200/60000 | epoch 4 | loss 0.0135 | grad 0.10 | 1.74 it/s | ETA 1.25h
+step 52300/60000 | epoch 4 | loss 0.0127 | grad 0.05 | 1.74 it/s | ETA 1.23h
+step 52400/60000 | epoch 4 | loss 0.0132 | grad 0.08 | 1.74 it/s | ETA 1.21h
+step 52500/60000 | epoch 4 | loss 0.0140 | grad 0.05 | 1.74 it/s | ETA 1.20h
+step 52600/60000 | epoch 4 | loss 0.0136 | grad 0.08 | 1.74 it/s | ETA 1.18h
+step 52700/60000 | epoch 4 | loss 0.0138 | grad 0.05 | 1.74 it/s | ETA 1.17h
+step 52800/60000 | epoch 4 | loss 0.0134 | grad 0.07 | 1.74 it/s | ETA 1.15h
+step 52900/60000 | epoch 4 | loss 0.0139 | grad 0.02 | 1.74 it/s | ETA 1.13h
+step 53000/60000 | epoch 4 | loss 0.0129 | grad 0.04 | 1.74 it/s | ETA 1.12h
+step 53100/60000 | epoch 4 | loss 0.0132 | grad 0.08 | 1.74 it/s | ETA 1.10h
+step 53200/60000 | epoch 4 | loss 0.0130 | grad 0.09 | 1.74 it/s | ETA 1.09h
+step 53300/60000 | epoch 4 | loss 0.0121 | grad 0.05 | 1.74 it/s | ETA 1.07h
+step 53400/60000 | epoch 4 | loss 0.0126 | grad 0.10 | 1.74 it/s | ETA 1.05h
+step 53500/60000 | epoch 4 | loss 0.0135 | grad 0.07 | 1.74 it/s | ETA 1.04h
+step 53600/60000 | epoch 4 | loss 0.0124 | grad 0.09 | 1.74 it/s | ETA 1.02h
+step 53700/60000 | epoch 4 | loss 0.0129 | grad 0.02 | 1.74 it/s | ETA 1.01h
+step 53800/60000 | epoch 4 | loss 0.0131 | grad 0.09 | 1.74 it/s | ETA 0.99h
+step 53900/60000 | epoch 4 | loss 0.0126 | grad 0.02 | 1.74 it/s | ETA 0.97h
+step 54000/60000 | epoch 4 | loss 0.0149 | grad 0.02 | 1.74 it/s | ETA 0.96h
+step 54100/60000 | epoch 4 | loss 0.0129 | grad 0.04 | 1.74 it/s | ETA 0.94h
+step 54200/60000 | epoch 4 | loss 0.0155 | grad 0.02 | 1.74 it/s | ETA 0.93h
+step 54300/60000 | epoch 4 | loss 0.0131 | grad 0.07 | 1.74 it/s | ETA 0.91h
+step 54400/60000 | epoch 4 | loss 0.0128 | grad 0.07 | 1.74 it/s | ETA 0.89h
+step 54500/60000 | epoch 4 | loss 0.0127 | grad 0.07 | 1.74 it/s | ETA 0.88h
+step 54600/60000 | epoch 4 | loss 0.0133 | grad 0.08 | 1.74 it/s | ETA 0.86h
+step 54700/60000 | epoch 4 | loss 0.0131 | grad 0.04 | 1.74 it/s | ETA 0.85h
+step 54800/60000 | epoch 4 | loss 0.0127 | grad 0.07 | 1.74 it/s | ETA 0.83h
+step 54900/60000 | epoch 4 | loss 0.0139 | grad 0.05 | 1.74 it/s | ETA 0.81h
+step 55000/60000 | epoch 4 | loss 0.0123 | grad 0.03 | 1.74 it/s | ETA 0.80h
+step 55100/60000 | epoch 4 | loss 0.0126 | grad 0.05 | 1.74 it/s | ETA 0.78h
+step 55200/60000 | epoch 4 | loss 0.0135 | grad 0.03 | 1.74 it/s | ETA 0.77h
+step 55300/60000 | epoch 4 | loss 0.0158 | grad 0.09 | 1.74 it/s | ETA 0.75h
+step 55400/60000 | epoch 4 | loss 0.0139 | grad 0.08 | 1.74 it/s | ETA 0.73h
+step 55500/60000 | epoch 4 | loss 0.0131 | grad 0.06 | 1.74 it/s | ETA 0.72h
+step 55600/60000 | epoch 4 | loss 0.0140 | grad 0.06 | 1.74 it/s | ETA 0.70h
+step 55700/60000 | epoch 4 | loss 0.0116 | grad 0.03 | 1.74 it/s | ETA 0.69h
+step 55800/60000 | epoch 4 | loss 0.0134 | grad 0.05 | 1.74 it/s | ETA 0.67h
+step 55900/60000 | epoch 4 | loss 0.0144 | grad 0.08 | 1.74 it/s | ETA 0.65h
+step 56000/60000 | epoch 4 | loss 0.0130 | grad 0.05 | 1.74 it/s | ETA 0.64h
+step 56100/60000 | epoch 4 | loss 0.0128 | grad 0.09 | 1.74 it/s | ETA 0.62h
+step 56200/60000 | epoch 4 | loss 0.0139 | grad 0.04 | 1.74 it/s | ETA 0.61h
+step 56300/60000 | epoch 4 | loss 0.0127 | grad 0.07 | 1.74 it/s | ETA 0.59h
+step 56400/60000 | epoch 4 | loss 0.0129 | grad 0.11 | 1.74 it/s | ETA 0.57h
+step 56500/60000 | epoch 4 | loss 0.0136 | grad 0.08 | 1.74 it/s | ETA 0.56h
+step 56600/60000 | epoch 4 | loss 0.0149 | grad 0.03 | 1.74 it/s | ETA 0.54h
+step 56700/60000 | epoch 4 | loss 0.0130 | grad 0.07 | 1.74 it/s | ETA 0.53h
+step 56800/60000 | epoch 4 | loss 0.0132 | grad 0.07 | 1.74 it/s | ETA 0.51h
+step 56900/60000 | epoch 4 | loss 0.0135 | grad 0.05 | 1.74 it/s | ETA 0.49h
+step 57000/60000 | epoch 4 | loss 0.0144 | grad 0.02 | 1.74 it/s | ETA 0.48h
+step 57100/60000 | epoch 4 | loss 0.0134 | grad 0.05 | 1.74 it/s | ETA 0.46h
+step 57200/60000 | epoch 4 | loss 0.0135 | grad 0.06 | 1.74 it/s | ETA 0.45h
+step 57300/60000 | epoch 4 | loss 0.0134 | grad 0.11 | 1.74 it/s | ETA 0.43h
+step 57400/60000 | epoch 4 | loss 0.0130 | grad 0.05 | 1.74 it/s | ETA 0.41h
+step 57500/60000 | epoch 4 | loss 0.0142 | grad 0.04 | 1.74 it/s | ETA 0.40h
+step 57600/60000 | epoch 4 | loss 0.0128 | grad 0.08 | 1.74 it/s | ETA 0.38h
+step 57700/60000 | epoch 4 | loss 0.0139 | grad 0.10 | 1.74 it/s | ETA 0.37h
+step 57800/60000 | epoch 4 | loss 0.0137 | grad 0.03 | 1.74 it/s | ETA 0.35h
+step 57900/60000 | epoch 4 | loss 0.0135 | grad 0.02 | 1.74 it/s | ETA 0.33h
+step 58000/60000 | epoch 4 | loss 0.0127 | grad 0.03 | 1.74 it/s | ETA 0.32h
+step 58100/60000 | epoch 4 | loss 0.0128 | grad 0.06 | 1.75 it/s | ETA 0.30h
+step 58200/60000 | epoch 4 | loss 0.0141 | grad 0.05 | 1.75 it/s | ETA 0.29h
+step 58300/60000 | epoch 4 | loss 0.0125 | grad 0.10 | 1.75 it/s | ETA 0.27h
+step 58400/60000 | epoch 4 | loss 0.0121 | grad 0.06 | 1.75 it/s | ETA 0.25h
+step 58500/60000 | epoch 4 | loss 0.0137 | grad 0.09 | 1.75 it/s | ETA 0.24h
+step 58600/60000 | epoch 4 | loss 0.0134 | grad 0.06 | 1.75 it/s | ETA 0.22h
+step 58700/60000 | epoch 4 | loss 0.0136 | grad 0.03 | 1.75 it/s | ETA 0.21h
+step 58800/60000 | epoch 4 | loss 0.0136 | grad 0.05 | 1.75 it/s | ETA 0.19h
+step 58900/60000 | epoch 4 | loss 0.0140 | grad 0.06 | 1.75 it/s | ETA 0.17h
+step 59000/60000 | epoch 4 | loss 0.0133 | grad 0.08 | 1.75 it/s | ETA 0.16h
+step 59100/60000 | epoch 4 | loss 0.0138 | grad 0.03 | 1.75 it/s | ETA 0.14h
+step 59200/60000 | epoch 4 | loss 0.0128 | grad 0.08 | 1.75 it/s | ETA 0.13h
+step 59300/60000 | epoch 4 | loss 0.0138 | grad 0.05 | 1.75 it/s | ETA 0.11h
+step 59400/60000 | epoch 4 | loss 0.0155 | grad 0.05 | 1.75 it/s | ETA 0.10h
+step 59500/60000 | epoch 4 | loss 0.0120 | grad 0.07 | 1.75 it/s | ETA 0.08h
+step 59600/60000 | epoch 4 | loss 0.0146 | grad 0.05 | 1.75 it/s | ETA 0.06h
+step 59700/60000 | epoch 4 | loss 0.0126 | grad 0.07 | 1.75 it/s | ETA 0.05h
+step 59800/60000 | epoch 4 | loss 0.0134 | grad 0.07 | 1.75 it/s | ETA 0.03h
+step 59900/60000 | epoch 4 | loss 0.0124 | grad 0.03 | 1.75 it/s | ETA 0.02h
+step 60000/60000 | epoch 4 | loss 0.0131 | grad 0.05 | 1.75 it/s | ETA 0.00h
+done: 9.54h | final pipeline at /media/joseph/Media/diffusion128/runs/celeba128/final

training/train.py

@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""
+Pretrain a small unconditional DDPM diffusion model at 128x128 from scratch
+on a single GPU (RTX 3060-friendly: fp16 AMP, ~50M-param UNet).
+
+Default dataset: tglcourse/CelebA-faces-cropped-128 (~202k SFW 128x128 face crops).
+Alternative:     huggan/flowers-102-categories
+
+Training stops at min(--max_train_steps, --max_hours wall clock), so it is
+safe to leave running overnight.
+
+Usage:
+    python train.py --output_dir runs/celeba128
+
+Sample from the finished run:
+    from diffusers import DiffusionPipeline
+    pipe = DiffusionPipeline.from_pretrained("runs/celeba128/final").to("cuda")
+    pipe(batch_size=8).images[0].save("sample.png")
+"""
+
+import argparse
+import copy
+import math
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from torchvision.utils import make_grid, save_image
+
+from datasets import load_dataset
+from diffusers import DDIMScheduler, DDPMPipeline, DDPMScheduler, UNet2DModel
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dataset_name", default="tglcourse/CelebA-faces-cropped-128")
+    p.add_argument("--split", default="train")
+    p.add_argument("--image_column", default="image")
+    p.add_argument("--resolution", type=int, default=128)
+    p.add_argument("--output_dir", default="runs/celeba128")
+    p.add_argument("--train_batch_size", type=int, default=16)
+    p.add_argument("--gradient_accumulation_steps", type=int, default=1)
+    p.add_argument("--max_train_steps", type=int, default=60_000)
+    p.add_argument("--max_hours", type=float, default=13.5,
+                   help="hard wall-clock budget; stops at steps OR hours, whichever first")
+    p.add_argument("--learning_rate", type=float, default=1e-4)
+    p.add_argument("--warmup_steps", type=int, default=500)
+    p.add_argument("--ema_decay", type=float, default=0.9999)
+    p.add_argument("--num_train_timesteps", type=int, default=1000)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--dataloader_num_workers", type=int, default=4)
+    p.add_argument("--log_every", type=int, default=100)
+    p.add_argument("--sample_every", type=int, default=3000)
+    p.add_argument("--num_sample_images", type=int, default=16)
+    p.add_argument("--ckpt_every", type=int, default=20_000)
+    p.add_argument("--keep_ckpts", type=int, default=2)
+    p.add_argument("--resume_from", default=None)
+    p.add_argument("--channels", type=int, nargs="+", default=[128, 128, 256, 256, 512])
+    p.add_argument("--layers_per_block", type=int, default=2)
+    p.add_argument("--no_tensorboard", action="store_true")
+    return p.parse_args()
+
+
+def build_unet(resolution, channels, layers_per_block):
+    n = len(channels)
+    attn_idx = n - 2
+    down = ["DownBlock2D"] * n
+    up = ["UpBlock2D"] * n
+    down[attn_idx] = "AttnDownBlock2D"
+    up[n - 1 - attn_idx] = "AttnUpBlock2D"
+    return UNet2DModel(
+        sample_size=resolution,
+        in_channels=3,
+        out_channels=3,
+        layers_per_block=layers_per_block,
+        block_out_channels=tuple(channels),
+        down_block_types=tuple(down),
+        up_block_types=tuple(up),
+    )
+
+
+def build_dataset(args):
+    ds = load_dataset(args.dataset_name, split=args.split)
+    tf = transforms.Compose([
+        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
+        transforms.CenterCrop(args.resolution),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5] * 3, [0.5] * 3),
+    ])
+
+    def apply(examples):
+        return {"images": [tf(img.convert("RGB")) for img in examples[args.image_column]]}
+
+    ds.set_transform(apply)
+    return ds
+
+
+@torch.no_grad()
+def ema_update(ema_model, model, decay):
+    for pe, pm in zip(ema_model.parameters(), model.parameters()):
+        pe.lerp_(pm.data, 1.0 - decay)
+
+
+@torch.no_grad()
+def generate_samples(ema_model, args, device, global_step):
+    ema_model.eval()
+    ddim = DDIMScheduler(num_train_timesteps=args.num_train_timesteps, beta_schedule="linear")
+    ddim.set_timesteps(50)
+    g = torch.Generator(device=device).manual_seed(args.seed)
+    x = torch.randn(args.num_sample_images, 3, args.resolution, args.resolution,
+                    generator=g, device=device, dtype=torch.float16)
+    for t in ddim.timesteps:
+        tt = t.expand(x.shape[0]).to(device)
+        with torch.autocast("cuda", dtype=torch.float16):
+            pred = ema_model(x, tt).sample
+        x = ddim.step(pred, t, x).prev_sample
+    x = (x.float().clamp(-1, 1) + 1.0) / 2.0
+    out = Path(args.output_dir) / "samples"
+    out.mkdir(parents=True, exist_ok=True)
+    save_image(make_grid(x, nrow=int(math.sqrt(x.shape[0]))), out / f"step_{global_step:07d}.png")
+    ema_model.train()
+
+
+def save_checkpoint(path, step, model, ema_model, optimizer, scaler):
+    torch.save({
+        "step": step,
+        "model": model.state_dict(),
+        "ema": ema_model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "scaler": scaler.state_dict(),
+    }, path)
+
+
+def prune_checkpoints(ckpt_dir, keep):
+    ckpts = sorted(ckpt_dir.glob("step_*.pt"))
+    for old in ckpts[:-keep]:
+        old.unlink()
+
+
+def main():
+    args = parse_args()
+    out = Path(args.output_dir)
+    ckpt_dir = out / "checkpoints"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+    torch.manual_seed(args.seed)
+    device = "cuda"
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+
+    writer = None
+    if not args.no_tensorboard:
+        try:
+            from torch.utils.tensorboard import SummaryWriter
+            writer = SummaryWriter(out / "tb")
+        except ImportError:
+            pass
+
+    ds = build_dataset(args)
+    loader = DataLoader(ds, batch_size=args.train_batch_size, shuffle=True,
+                        num_workers=args.dataloader_num_workers, pin_memory=True,
+                        drop_last=True)
+
+    model = build_unet(args.resolution, args.channels, args.layers_per_block).to(device)
+    ema_model = copy.deepcopy(model).eval()
+    for p in ema_model.parameters():
+        p.requires_grad_(False)
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"model params: {n_params / 1e6:.1f}M | dataset rows: {len(ds)} | "
+          f"batch {args.train_batch_size} x accum {args.gradient_accumulation_steps}", flush=True)
+
+    T = args.num_train_timesteps
+    noise_sched = DDPMScheduler(num_train_timesteps=T, beta_schedule="linear",
+                                prediction_type="epsilon")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
+    scaler = torch.amp.GradScaler("cuda")
+
+    start_step = 0
+    if args.resume_from:
+        st = torch.load(args.resume_from, map_location=device)
+        model.load_state_dict(st["model"])
+        ema_model.load_state_dict(st["ema"])
+        optimizer.load_state_dict(st["optimizer"])
+        scaler.load_state_dict(st["scaler"])
+        start_step = st["step"]
+        print(f"resumed from {args.resume_from} at step {start_step}", flush=True)
+
+    lr_sched = torch.optim.lr_scheduler.LambdaLR(
+        optimizer, lambda s: min(1.0, (s + 1) / max(1, args.warmup_steps)),
+        last_epoch=start_step - 1)
+
+    accum = args.gradient_accumulation_steps
+    model.train()
+    optimizer.zero_grad(set_to_none=True)
+    data_iter = iter(loader)
+    epoch = 0
+    t0 = time.time()
+    budget_s = args.max_hours * 3600
+    running = 0.0
+
+    for step in range(start_step, args.max_train_steps):
+        if time.time() - t0 >= budget_s:
+            print(f"wall-clock budget of {args.max_hours}h reached at step {step}", flush=True)
+            break
+        try:
+            batch = next(data_iter)
+        except StopIteration:
+            epoch += 1
+            data_iter = iter(loader)
+            batch = next(data_iter)
+
+        x0 = batch["images"].to(device, non_blocking=True)
+        noise = torch.randn_like(x0)
+        timesteps = torch.randint(0, T, (x0.shape[0],), device=device, dtype=torch.long)
+        noisy = noise_sched.add_noise(x0, noise, timesteps)
+
+        with torch.autocast("cuda", dtype=torch.float16):
+            pred = model(noisy, timesteps).sample
+            loss = F.mse_loss(pred.float(), noise)
+        scaler.scale(loss / accum).backward()
+
+        if (step + 1) % accum == 0:
+            scaler.unscale_(optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+            lr_sched.step()
+            ema_update(ema_model, model, args.ema_decay)
+        running += loss.item()
+
+        if writer and (step + 1) % args.log_every == 0:
+            writer.add_scalar("train/loss", running / args.log_every, step + 1)
+            writer.add_scalar("train/lr", lr_sched.get_last_lr()[0], step + 1)
+
+        if (step + 1) % args.log_every == 0:
+            el = time.time() - t0
+            sps = (step + 1 - start_step) / max(el, 1e-6)
+            eta = min((args.max_train_steps - (step + 1)) / max(sps, 1e-6), budget_s - el)
+            print(f"step {step + 1}/{args.max_train_steps} | epoch {epoch} | "
+                  f"loss {running / args.log_every:.4f} | grad {grad_norm:.2f} | "
+                  f"{sps:.2f} it/s | ETA {max(eta, 0) / 3600:.2f}h", flush=True)
+            running = 0.0
+
+        if (step + 1) % args.sample_every == 0:
+            generate_samples(ema_model, args, device, step + 1)
+
+        if (step + 1) % args.ckpt_every == 0:
+            save_checkpoint(ckpt_dir / f"step_{step + 1}.pt", step + 1,
+                            model, ema_model, optimizer, scaler)
+            prune_checkpoints(ckpt_dir, args.keep_ckpts)
+
+    save_checkpoint(ckpt_dir / "last.pt", step + 1, model, ema_model, optimizer, scaler)
+    generate_samples(ema_model, args, device, step + 1)
+    pipe = DDPMPipeline(unet=ema_model,
+                        scheduler=DDPMScheduler(num_train_timesteps=T, beta_schedule="linear"))
+    pipe.save_pretrained(out / "final")
+    if writer:
+        writer.close()
+    print(f"done: {(time.time() - t0) / 3600:.2f}h | final pipeline at {out / 'final'}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

unet/config.json

@@ -0,0 +1,50 @@
+{
+  "_class_name": "UNet2DModel",
+  "_diffusers_version": "0.39.0",
+  "act_fn": "silu",
+  "add_attention": true,
+  "attention_head_dim": 8,
+  "attn_norm_num_groups": null,
+  "block_out_channels": [
+    128,
+    128,
+    256,
+    256,
+    512
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "down_block_types": [
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D",
+    "AttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "downsample_type": "conv",
+  "dropout": 0.0,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 3,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2D",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_class_embeds": null,
+  "num_train_timesteps": null,
+  "out_channels": 3,
+  "resnet_time_scale_shift": "default",
+  "sample_size": 128,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "up_block_types": [
+    "UpBlock2D",
+    "AttnUpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D"
+  ],
+  "upsample_type": "conv"
+}

unet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3230ecffd2276430652df93f7ab970846a3537f6b54e1c117487e628093665c5
+size 285788892