Struct Trainer

Source

pub struct Trainer {Show 20 fields
    pub environment: Environment,
    pub training_images_repeat: usize,
    pub regularization_images_repeat: usize,
    pub resolution: (usize, usize),
    pub save_model_as: ModelFileFormat,
    pub network_module: String,
    pub text_encoder_lr: f32,
    pub unet_lr: f32,
    pub lr_scheduler_num_cycles: usize,
    pub learning_rate: f32,
    pub lr_warmup_steps: usize,
    pub train_batch_size: usize,
    pub max_train_steps: usize,
    pub save_every_n_epochs: usize,
    pub mixed_precision: FloatPrecision,
    pub save_precision: FloatPrecision,
    pub max_grad_norm: f32,
    pub max_data_loader_n_workers: usize,
    pub bucket_reso_steps: usize,
    pub noise_offset: f32,
}

Expand description

The Trainer structure.

Fields§

§environment: Environment

The environment to use for the training process.

§training_images_repeat: usize

The number of times to repeat the training images.

§regularization_images_repeat: usize

The number of times to repeat the regularization images.

§resolution: (usize, usize)

The maximum resolution of the images to use for the training process.

§save_model_as: ModelFileFormat

The format to save the model as.

§network_module: String

The module to use for the network.

§text_encoder_lr: f32

The learning rate for the text encoder.

§unet_lr: f32

The learning rate for the unet.

§lr_scheduler_num_cycles: usize

The number of cycles for the learning rate scheduler.

§learning_rate: f32

The learning rate for the training process.

§lr_warmup_steps: usize

The number of warmup steps for the learning rate.

§train_batch_size: usize

The batch size for the training process.

§max_train_steps: usize

The maximum number of training steps.

§save_every_n_epochs: usize

The frequency to

§mixed_precision: FloatPrecision

The precision to use for mixed precision training.

§save_precision: FloatPrecision

The precision to use for saving the model.

§max_grad_norm: f32

The maximum gradient norm.

§max_data_loader_n_workers: usize

The maximum number of data loader workers.

§bucket_reso_steps: usize

The number of steps for the bucket resolution.

§noise_offset: f32

The noise offset.

Implementations§

Source §

impl Trainer

Source

pub fn new() -> Self

Create a new Trainer.

Examples found in repository ?

examples/train-lora.rs (line 13)

3fn main() {
4    let kohya_ss = std::env::var("KOHYA_SS_PATH").expect("KOHYA_SS_PATH not set");
5    let environment = Environment::new().with_kohya_ss(kohya_ss);
6
7    let prompt = Prompt::new("bacana", "white dog");
8    let image_data_set = ImageDataSet::from_dir("examples/training/lora/bacana/images");
9    let data_set = TrainingDataSet::new(image_data_set);
10    let output = Output::new("{prompt.instance}({prompt.class})d{network.dimension}a{network.alpha}", "examples/training/lora/bacana/output");
11    let parameters = Parameters::new(prompt, data_set, output);
12
13    Trainer::new()
14        .with_environment(environment)
15        .start(&parameters);
16}

Source

pub fn with_environment(self, environment: Environment) -> Self

Set the environment for the training process.

Examples found in repository ?

examples/train-lora.rs (line 14)

3fn main() {
4    let kohya_ss = std::env::var("KOHYA_SS_PATH").expect("KOHYA_SS_PATH not set");
5    let environment = Environment::new().with_kohya_ss(kohya_ss);
6
7    let prompt = Prompt::new("bacana", "white dog");
8    let image_data_set = ImageDataSet::from_dir("examples/training/lora/bacana/images");
9    let data_set = TrainingDataSet::new(image_data_set);
10    let output = Output::new("{prompt.instance}({prompt.class})d{network.dimension}a{network.alpha}", "examples/training/lora/bacana/output");
11    let parameters = Parameters::new(prompt, data_set, output);
12
13    Trainer::new()
14        .with_environment(environment)
15        .start(&parameters);
16}

Source

pub fn start(&mut self, parameters: &Parameters)

Start the training process.

Examples found in repository ?

examples/train-lora.rs (line 15)

3fn main() {
4    let kohya_ss = std::env::var("KOHYA_SS_PATH").expect("KOHYA_SS_PATH not set");
5    let environment = Environment::new().with_kohya_ss(kohya_ss);
6
7    let prompt = Prompt::new("bacana", "white dog");
8    let image_data_set = ImageDataSet::from_dir("examples/training/lora/bacana/images");
9    let data_set = TrainingDataSet::new(image_data_set);
10    let output = Output::new("{prompt.instance}({prompt.class})d{network.dimension}a{network.alpha}", "examples/training/lora/bacana/output");
11    let parameters = Parameters::new(prompt, data_set, output);
12
13    Trainer::new()
14        .with_environment(environment)
15        .start(&parameters);
16}