Repository navigation
pixtreme-train v0.3.0
Training from a YAML configuration:
registerandlookupselect networks, losses, training steps, optimizers, schedulers and datasets by name. Thedefaultstep, theadam,adamwandsgdoptimizers, thewarmup_cosinescheduler, thel1,mseandcharbonnierlosses, and theimage_degradationandimage_evaluationdatasets are registered.TrainingStepis the base class for one update; a trainer overridesstepand, when needed, the optimizer construction, state, configuration checks and evaluation.DefaultSteptrains onemainnetwork 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. trainruns 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 runstrainand prints the path of the last checkpoint. load_networkrebuilds a trained network from a checkpoint, with its EMA or training weights.
Image data and on-the-fly LQ generation:
DegradationPipelineandload_pipelinedescribe named, seeded sequences ofresizeandencodesteps (jpeg,prores_lt,h264,hevc) with fixed values, ranges or candidates.ImageDegradationDatasetcrops aligned GT / LQ patches with randomness derived from the sample number, so samples do not depend on DataLoader workers;NumberedSamplerandsample_generatorprovide those numbers and streams.augment_imagesflips and transposes paired images.ImageEvaluationDatasetdegrades full evaluation images;image_pathsselects 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.