Skip to content

v0.3.0

Latest

Choose a tag to compare

@minamikik minamikik released this 28 Oct 22:40

pixtreme-train v0.3.0

Training from a YAML configuration:

  • register and lookup select networks, losses, training steps, optimizers, schedulers and datasets by name. The default step, the adam, adamw and sgd optimizers, the warmup_cosine scheduler, the l1, mse and charbonnier losses, and the image_degradation and image_evaluation datasets are registered.
  • TrainingStep is the base class for one update; a trainer overrides step and, when needed, the optimizer construction, state, configuration checks and evaluation. DefaultStep trains one main network with weighted losses, a scheduler and optional EMA.
  • With train.precision: bf16, only the network forward runs under bfloat16 autocast on CUDA; the loss, weights, gradients, optimizer state, EMA weights, evaluation and checkpoints stay float32.
  • train runs the training loop: resume from a checkpoint, checkpoint and evaluation intervals, one JSON line per update, and SIGINT / SIGTERM handling that saves a checkpoint before exiting.
  • The pixtreme-train train CONFIG --output DIR [--resume CHECKPOINT] command runs train and prints the path of the last checkpoint.
  • load_network rebuilds a trained network from a checkpoint, with its EMA or training weights.

Image data and on-the-fly LQ generation:

  • DegradationPipeline and load_pipeline describe named, seeded sequences of resize and encode steps (jpeg, prores_lt, h264, hevc) with fixed values, ranges or candidates.
  • ImageDegradationDataset crops aligned GT / LQ patches with randomness derived from the sample number, so samples do not depend on DataLoader workers; NumberedSampler and sample_generator provide those numbers and streams. augment_images flips and transposes paired images.
  • ImageEvaluationDataset degrades full evaluation images; image_paths selects image files.

The prores_lt, h264 and hevc steps, and reading 16-bit RGB images, need an ffmpeg command on PATH. Importing pixtreme.train.common still does not load torch.

The wheel attached here is the reference copy for the PyPI distribution of pixtreme-train. Verify it with sha256sum -c SHA256SUMS.