Added validation loss
This commit is contained in:
Jaret Burkett
parent
cd677c70b5
commit
1eb97b7443
|
|
@ -12,6 +12,8 @@ from typing import Union, List, Optional
|
|||
|
||||
import numpy as np
|
||||
import yaml
|
||||
from PIL import Image, ImageOps
|
||||
from torchvision import transforms
|
||||
from diffusers import T2IAdapter, ControlNetModel
|
||||
from diffusers.training_utils import compute_density_for_timestep_sampling
|
||||
from safetensors.torch import save_file, load_file
|
||||
|
|
@ -40,6 +42,7 @@ from toolkit.network_mixins import Network
|
|||
from toolkit.optimizer import get_optimizer
|
||||
from toolkit.paths import CONFIG_ROOT
|
||||
from toolkit.progress_bar import ToolkitProgressBar
|
||||
from toolkit.prompt_utils import concat_prompt_embeds
|
||||
from toolkit.reference_adapter import ReferenceAdapter
|
||||
from toolkit.sampler import get_sampler
|
||||
from toolkit.saving import save_t2i_from_diffusers, load_t2i_model, save_ip_adapter_from_diffusers, \
|
||||
|
|
@ -47,7 +50,7 @@ from toolkit.saving import save_t2i_from_diffusers, load_t2i_model, save_ip_adap
|
|||
|
||||
from toolkit.scheduler import get_lr_scheduler
|
||||
from toolkit.sd_device_states_presets import get_train_sd_device_state_preset
|
||||
from toolkit.stable_diffusion_model import StableDiffusion
|
||||
from toolkit.stable_diffusion_model import StableDiffusion, BlankNetwork
|
||||
|
||||
from jobs.process import BaseTrainProcess
|
||||
from toolkit.metadata import get_meta_for_safetensors, load_metadata_from_safetensors, add_base_model_info_to_meta, \
|
||||
|
|
@ -261,6 +264,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
|
|||
self.steps_this_boundary = 0
|
||||
self.num_consecutive_oom = 0
|
||||
self.additional_logs = {}
|
||||
# cached latents, prompt embeds, and fixed noise for validation
|
||||
self._validation_cache = None
|
||||
|
||||
def post_process_generate_image_config_list(self, generate_image_config_list: List[GenerateImageConfig]):
|
||||
# override in subclass
|
||||
|
|
@ -1570,6 +1575,146 @@ class BaseSDTrainProcess(BaseTrainProcess):
|
|||
self.load_training_state_from_metadata(latest_save_path)
|
||||
# set trainable params
|
||||
self.sd.adapter = self.adapter
|
||||
|
||||
def setup_validation(self):
|
||||
# caches everything needed for validation (latents, prompt embeds, fixed noise)
|
||||
# must be called while the vae and text encoder are still loaded, they may be
|
||||
# dumped later to save memory
|
||||
val_config = self.train_config.validation_config
|
||||
if val_config is None:
|
||||
return
|
||||
if not self.accelerator.is_main_process:
|
||||
return
|
||||
validation_items = []
|
||||
for item in val_config.validation_items:
|
||||
if not item.image_path:
|
||||
print_acc("Skipping validation item with no image")
|
||||
continue
|
||||
if not os.path.exists(item.image_path):
|
||||
print_acc(f"Skipping validation item, image not found: {item.image_path}")
|
||||
continue
|
||||
validation_items.append(item)
|
||||
if len(validation_items) == 0:
|
||||
print_acc("Validation config has no valid validation_items, skipping validation")
|
||||
return
|
||||
print_acc(f"Caching validation latents and embeddings for {len(validation_items)} images")
|
||||
device = self.device_torch
|
||||
dtype = get_torch_dtype(self.train_config.dtype)
|
||||
resolution = val_config.resolution
|
||||
|
||||
image_list = []
|
||||
prompt_list = []
|
||||
for item in validation_items:
|
||||
img = Image.open(item.image_path)
|
||||
img = ImageOps.exif_transpose(img).convert('RGB')
|
||||
# deterministic resize, shortest side to resolution then center crop
|
||||
img = transforms.Resize(resolution)(img)
|
||||
img = transforms.CenterCrop(resolution)(img)
|
||||
tensor = transforms.ToTensor()(img) * 2.0 - 1.0
|
||||
image_list.append(tensor)
|
||||
prompt_list.append(item.prompt)
|
||||
|
||||
fork_devices = [device] if device.type == 'cuda' else []
|
||||
with torch.no_grad(), torch.random.fork_rng(devices=fork_devices):
|
||||
# encode the prompts one at a time so they can be reassembled per sigma later
|
||||
te = self.sd.text_encoder
|
||||
te_list = te if isinstance(te, list) else ([te] if te is not None else [])
|
||||
orig_te_devices = [next(t.parameters()).device for t in te_list]
|
||||
self.sd.text_encoder_to(device)
|
||||
embeds_list = [
|
||||
self.sd.encode_prompt([prompt]).to('cpu', dtype=torch.float32).detach()
|
||||
for prompt in prompt_list
|
||||
]
|
||||
for t, te_device in zip(te_list, orig_te_devices):
|
||||
t.to(te_device)
|
||||
|
||||
# seed so the vae latent dist sampling is always identical
|
||||
torch.manual_seed(42)
|
||||
orig_vae_device = self.sd.vae.device
|
||||
latents = self.sd.encode_images(image_list, device=device, dtype=dtype)
|
||||
self.sd.vae.to(orig_vae_device)
|
||||
|
||||
# fixed noise per image, seeds start at 42 and increment for each image
|
||||
noise_list = []
|
||||
for i in range(latents.shape[0]):
|
||||
generator = torch.Generator(device='cpu').manual_seed(42 + i)
|
||||
noise_list.append(
|
||||
torch.randn(latents[i:i + 1].shape, generator=generator, dtype=torch.float32)
|
||||
)
|
||||
noise = torch.cat(noise_list, dim=0)
|
||||
|
||||
self._validation_cache = {
|
||||
'latents': latents.to('cpu', dtype=torch.float32),
|
||||
'noise': noise,
|
||||
'embeds': embeds_list,
|
||||
}
|
||||
flush()
|
||||
|
||||
def validate(self):
|
||||
val_config = self.train_config.validation_config
|
||||
if val_config is None or self._validation_cache is None:
|
||||
return
|
||||
if not self.accelerator.is_main_process:
|
||||
return
|
||||
device = self.device_torch
|
||||
dtype = get_torch_dtype(self.train_config.dtype)
|
||||
sigmas = val_config.validation_sigmas
|
||||
cache = self._validation_cache
|
||||
|
||||
was_unet_training = self.sd.unet.training
|
||||
self.sd.unet.eval()
|
||||
# the network is only active inside its context, without this the base model is validated
|
||||
network = self.network if self.network is not None else BlankNetwork()
|
||||
start_multiplier = network.multiplier
|
||||
network.multiplier = 1.0
|
||||
with torch.no_grad(), network:
|
||||
latents = cache['latents'].to(device, dtype=dtype)
|
||||
noise = cache['noise'].to(device, dtype=dtype)
|
||||
num_images = latents.shape[0]
|
||||
|
||||
# single batch of every (image, sigma) pair, ordered sigma major
|
||||
batch_latents = torch.cat([latents] * len(sigmas), dim=0)
|
||||
batch_noise = torch.cat([noise] * len(sigmas), dim=0)
|
||||
batch_embeds = concat_prompt_embeds(
|
||||
[e.clone().to(device, dtype=dtype) for e in cache['embeds']] * len(sigmas)
|
||||
)
|
||||
|
||||
if self.sd.is_flow_matching:
|
||||
timestep_values = [sigma * 1000.0 for sigma in sigmas]
|
||||
else:
|
||||
num_train_timesteps = self.sd.noise_scheduler.config.num_train_timesteps
|
||||
timestep_values = [
|
||||
min(int(round(sigma * num_train_timesteps)), num_train_timesteps - 1)
|
||||
for sigma in sigmas
|
||||
]
|
||||
timesteps = torch.tensor(
|
||||
[t for t in timestep_values for _ in range(num_images)],
|
||||
device=device,
|
||||
)
|
||||
|
||||
noisy_latents = self.sd.add_noise(batch_latents, batch_noise, timesteps).detach()
|
||||
|
||||
noise_pred = self.sd.predict_noise(
|
||||
latents=noisy_latents.to(device, dtype=dtype),
|
||||
conditional_embeddings=batch_embeds,
|
||||
timestep=timesteps,
|
||||
guidance_scale=1.0,
|
||||
guidance_embedding_scale=self.train_config.cfg_scale,
|
||||
bypass_guidance_embedding=self.train_config.bypass_guidance_embedding,
|
||||
)
|
||||
|
||||
if self.sd.is_flow_matching:
|
||||
target = batch_noise - batch_latents
|
||||
elif self.sd.prediction_type == 'v_prediction':
|
||||
target = self.sd.noise_scheduler.get_velocity(batch_latents, batch_noise, timesteps)
|
||||
else:
|
||||
target = batch_noise
|
||||
|
||||
val_loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float())
|
||||
self.additional_logs['val/loss'] = val_loss.item()
|
||||
network.multiplier = start_multiplier
|
||||
if was_unet_training:
|
||||
self.sd.unet.train()
|
||||
|
||||
def run(self):
|
||||
# torch.autograd.set_detect_anomaly(True)
|
||||
|
|
@ -2065,6 +2210,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
|
|||
)
|
||||
self.lr_scheduler = lr_scheduler
|
||||
|
||||
# cache validation latents and embeddings now, the vae and text encoder
|
||||
# may be dumped before the train loop starts
|
||||
self.setup_validation()
|
||||
|
||||
### HOOk ###
|
||||
self.before_dataset_load()
|
||||
# load datasets if passed in the root process
|
||||
|
|
@ -2480,6 +2629,18 @@ class BaseSDTrainProcess(BaseTrainProcess):
|
|||
print("\n==== Profile Results ====")
|
||||
print(self.torch_profiler.key_averages().table(sort_by="cpu_time_total", row_limit=1000))
|
||||
self.timer.stop('train_loop')
|
||||
|
||||
# run validation before any possible sampling/logging for this step
|
||||
if not did_oom and self.train_config.validation_config is not None:
|
||||
val_config = self.train_config.validation_config
|
||||
is_validate_step = (
|
||||
self.step_num == self.start_step
|
||||
or (val_config.validate_every_n_steps and self.step_num % val_config.validate_every_n_steps == 0)
|
||||
)
|
||||
if is_validate_step:
|
||||
with self.timer('validate'):
|
||||
self.validate()
|
||||
|
||||
if not did_first_flush:
|
||||
flush()
|
||||
did_first_flush = True
|
||||
|
|
|
|||
|
|
@ -337,6 +337,23 @@ class AdapterConfig:
|
|||
self.i2v_do_start_frame: bool = kwargs.get('i2v_do_start_frame', False)
|
||||
|
||||
|
||||
class ValidationItem:
|
||||
def __init__(self, **kwargs):
|
||||
self.image_path: str = kwargs.get('image_path', '')
|
||||
self.prompt: str = kwargs.get('prompt', '')
|
||||
|
||||
|
||||
class ValidationConfig:
|
||||
def __init__(self, **kwargs):
|
||||
self.validation_items: List[ValidationItem] = [
|
||||
item if isinstance(item, ValidationItem) else ValidationItem(**item)
|
||||
for item in kwargs.get('validation_items', [])
|
||||
]
|
||||
self.resolution: int = kwargs.get('resolution', 512)
|
||||
self.validate_every_n_steps: int = kwargs.get('validate_every_n_steps', 10)
|
||||
self.validation_sigmas: List[float] = kwargs.get('validation_sigmas', [1.0, 0.75, 0.5, 0.25])
|
||||
|
||||
|
||||
class EmbeddingConfig:
|
||||
def __init__(self, **kwargs):
|
||||
self.trigger = kwargs.get('trigger', 'custom_embedding')
|
||||
|
|
@ -592,6 +609,10 @@ class TrainConfig:
|
|||
self.max_loss_debug: bool = kwargs.get("max_loss_debug", False)
|
||||
# will clip the loss to this amount to prevent wild outliers
|
||||
self.max_loss: Optional[float] = kwargs.get("max_loss", None)
|
||||
self.validation_config: Optional[ValidationConfig] = None
|
||||
validation = kwargs.get('validation_config', None)
|
||||
if validation is not None:
|
||||
self.validation_config: ValidationConfig = ValidationConfig(**validation)
|
||||
|
||||
|
||||
ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21', 'anima']
|
||||
|
|
@ -1417,5 +1438,3 @@ def validate_configs(
|
|||
|
||||
if train_config.batch_size > 1 and any(dataset_config.auto_frame_count for dataset_config in dataset_configs):
|
||||
raise ValueError("Cannot use batch size greater than 1 with auto_frame_count. Please set batch_size to 1 or auto_frame_count to False.")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -199,6 +199,8 @@ export default function SimpleJob({
|
|||
|
||||
const showGPUSelect = !isMac();
|
||||
|
||||
const validationConfig = jobConfig.config.process[0].train.validation_config;
|
||||
|
||||
let numDatasetCols = 4;
|
||||
let numSampleTopCols = 4;
|
||||
let datasetStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6';
|
||||
|
|
@ -834,6 +836,136 @@ export default function SimpleJob({
|
|||
</div>
|
||||
</Card>
|
||||
</div>
|
||||
<div>
|
||||
<Card
|
||||
title="Validation"
|
||||
toggled={!!validationConfig}
|
||||
onToggle={value => {
|
||||
if (value) {
|
||||
setJobConfig(
|
||||
{
|
||||
validation_items: [{ image_path: '', prompt: '' }],
|
||||
resolution: 512,
|
||||
validate_every_n_steps: 10,
|
||||
validation_sigmas: [1.0, 0.75, 0.5, 0.25],
|
||||
},
|
||||
'config.process[0].train.validation_config',
|
||||
);
|
||||
} else {
|
||||
setJobConfig(undefined, 'config.process[0].train.validation_config');
|
||||
}
|
||||
}}
|
||||
>
|
||||
{validationConfig && (
|
||||
<>
|
||||
<p className="text-sm text-gray-400 mb-4">
|
||||
Validation runs a stable loss check on a fixed set of images. Each image is encoded once at startup
|
||||
and predicted at the selected sigmas with fixed seeds, so the result is always deterministic and
|
||||
comparable across the run. The average loss is logged as val/loss every time validation runs. The
|
||||
images need to match the concept of your dataset, but{' '}
|
||||
<span className="font-bold text-gray-300">do not include the validation images in the dataset</span>.
|
||||
They must be images containing the concept you want to train, but not an image trained on.
|
||||
</p>
|
||||
<div className="grid grid-cols-1 md:grid-cols-3 gap-6">
|
||||
<NumberInput
|
||||
label="Validate Every"
|
||||
value={validationConfig.validate_every_n_steps}
|
||||
onChange={value =>
|
||||
setJobConfig(value, 'config.process[0].train.validation_config.validate_every_n_steps')
|
||||
}
|
||||
placeholder="eg. 10"
|
||||
min={1}
|
||||
required
|
||||
/>
|
||||
<NumberInput
|
||||
label="Validation Resolution"
|
||||
value={validationConfig.resolution}
|
||||
onChange={value => setJobConfig(value, 'config.process[0].train.validation_config.resolution')}
|
||||
placeholder="eg. 512"
|
||||
min={64}
|
||||
required
|
||||
/>
|
||||
<SelectInput
|
||||
label="Validation Sigmas"
|
||||
value={(validationConfig.validation_sigmas ?? [1.0, 0.75, 0.5, 0.25]).join(', ')}
|
||||
onChange={value =>
|
||||
setJobConfig(
|
||||
value.split(',').map((v: string) => parseFloat(v)),
|
||||
'config.process[0].train.validation_config.validation_sigmas',
|
||||
)
|
||||
}
|
||||
options={[
|
||||
{ value: '0.5', label: '0.5' },
|
||||
{ value: '1, 0.5', label: '1.0, 0.5' },
|
||||
{ value: '1, 0.66, 0.33', label: '1.0, 0.66, 0.33' },
|
||||
{ value: '1, 0.75, 0.5, 0.25', label: '1.0, 0.75, 0.5, 0.25' },
|
||||
]}
|
||||
/>
|
||||
</div>
|
||||
<div className="mt-4">
|
||||
<label className="block text-xs text-gray-300 mb-2">
|
||||
Validation Images ({validationConfig.validation_items.length})
|
||||
</label>
|
||||
{validationConfig.validation_items.map((item, i) => (
|
||||
<div key={i} className="rounded-lg pl-4 pr-1 py-3 mb-4 bg-gray-950">
|
||||
<div className="flex items-center space-x-4">
|
||||
<SampleControlImage
|
||||
instruction="Add Image"
|
||||
src={item.image_path === '' ? null : item.image_path}
|
||||
onNewImageSelected={imagePath => {
|
||||
setJobConfig(
|
||||
imagePath ?? '',
|
||||
`config.process[0].train.validation_config.validation_items[${i}].image_path`,
|
||||
);
|
||||
}}
|
||||
/>
|
||||
<div className="flex-1">
|
||||
<TextInput
|
||||
label="Prompt"
|
||||
value={item.prompt}
|
||||
onChange={value =>
|
||||
setJobConfig(
|
||||
value,
|
||||
`config.process[0].train.validation_config.validation_items[${i}].prompt`,
|
||||
)
|
||||
}
|
||||
placeholder="Enter prompt"
|
||||
/>
|
||||
</div>
|
||||
<div>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() =>
|
||||
setJobConfig(
|
||||
validationConfig.validation_items.filter((_, index) => index !== i),
|
||||
'config.process[0].train.validation_config.validation_items',
|
||||
)
|
||||
}
|
||||
className="rounded-full p-1 text-sm"
|
||||
>
|
||||
<X />
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() =>
|
||||
setJobConfig(
|
||||
[...validationConfig.validation_items, { image_path: '', prompt: '' }],
|
||||
'config.process[0].train.validation_config.validation_items',
|
||||
)
|
||||
}
|
||||
className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors"
|
||||
>
|
||||
Add Validation Image
|
||||
</button>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</Card>
|
||||
</div>
|
||||
<div>
|
||||
<Card title="Advanced" collapsible>
|
||||
<div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
|
||||
|
|
|
|||
|
|
@ -7,9 +7,52 @@ interface CardProps {
|
|||
children?: React.ReactNode;
|
||||
collapsible?: boolean;
|
||||
defaultOpen?: boolean;
|
||||
// when provided, the card is opened/closed by a toggle switch on the right instead of a chevron
|
||||
toggled?: boolean;
|
||||
onToggle?: (value: boolean) => void;
|
||||
}
|
||||
|
||||
const Card: React.FC<CardProps> = ({ title, children, collapsible, defaultOpen }) => {
|
||||
const Card: React.FC<CardProps> = ({ title, children, collapsible, defaultOpen, toggled, onToggle }) => {
|
||||
if (onToggle) {
|
||||
return (
|
||||
<section className="space-y-2 px-4 pb-2 pt-2 bg-gray-900 rounded-lg">
|
||||
<div className="flex items-center justify-between">
|
||||
<div className="flex-1">
|
||||
{title && (
|
||||
<h2 className={classNames('text-lg font-semibold uppercase text-gray-500', toggled ? 'mb-2' : 'mb-0')}>
|
||||
{title}
|
||||
</h2>
|
||||
)}
|
||||
</div>
|
||||
<button
|
||||
type="button"
|
||||
role="switch"
|
||||
aria-checked={toggled}
|
||||
onClick={() => onToggle(!toggled)}
|
||||
className={classNames(
|
||||
'relative ml-2 inline-flex h-6 w-11 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-blue-600 focus:ring-offset-2',
|
||||
toggled ? 'bg-blue-500' : 'bg-gray-600',
|
||||
'hover:bg-opacity-80',
|
||||
)}
|
||||
>
|
||||
<span className="sr-only">Toggle {title}</span>
|
||||
<span
|
||||
className={classNames(
|
||||
'pointer-events-none inline-block h-5 w-5 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
|
||||
toggled ? 'translate-x-5' : 'translate-x-0',
|
||||
)}
|
||||
/>
|
||||
</button>
|
||||
</div>
|
||||
{toggled && (
|
||||
<>
|
||||
{children ?? null}
|
||||
<div className="pt-2"></div>
|
||||
</>
|
||||
)}
|
||||
</section>
|
||||
);
|
||||
}
|
||||
if (collapsible) {
|
||||
return (
|
||||
<Disclosure as="section" className="space-y-2 px-4 pb-2 pt-2 bg-gray-900 rounded-lg" defaultOpen={defaultOpen}>
|
||||
|
|
|
|||
|
|
@ -198,18 +198,26 @@ export const NumberInput = (props: NumberInputProps) => {
|
|||
|
||||
const numValue = Number(rawValue);
|
||||
|
||||
// Only apply constraints and call onChange when we have a valid number
|
||||
// don't clamp to min/max while typing, it mangles partial input (typing 1024 with
|
||||
// min 64 becomes 64024). Clamping happens on blur.
|
||||
if (!isNaN(numValue)) {
|
||||
let constrainedValue = numValue;
|
||||
|
||||
// Apply min/max constraints if they exist
|
||||
if (min !== undefined && constrainedValue < min) {
|
||||
constrainedValue = min;
|
||||
}
|
||||
if (max !== undefined && constrainedValue > max) {
|
||||
constrainedValue = max;
|
||||
}
|
||||
|
||||
onChange(numValue);
|
||||
}
|
||||
}}
|
||||
onBlur={() => {
|
||||
const numValue = Number(inputValue);
|
||||
if (inputValue === '' || isNaN(numValue)) {
|
||||
return;
|
||||
}
|
||||
let constrainedValue = numValue;
|
||||
if (min !== undefined && constrainedValue < min) {
|
||||
constrainedValue = min;
|
||||
}
|
||||
if (max !== undefined && constrainedValue > max) {
|
||||
constrainedValue = max;
|
||||
}
|
||||
if (constrainedValue !== numValue) {
|
||||
setInputValue(constrainedValue);
|
||||
onChange(constrainedValue);
|
||||
}
|
||||
}}
|
||||
|
|
|
|||
|
|
@ -117,6 +117,18 @@ export interface EMAConfig {
|
|||
ema_decay: number;
|
||||
}
|
||||
|
||||
export interface ValidationItem {
|
||||
image_path: string;
|
||||
prompt: string;
|
||||
}
|
||||
|
||||
export interface ValidationConfig {
|
||||
validation_items: ValidationItem[];
|
||||
resolution: number;
|
||||
validate_every_n_steps: number;
|
||||
validation_sigmas?: number[];
|
||||
}
|
||||
|
||||
export interface TrainConfig {
|
||||
batch_size: number;
|
||||
bypass_guidance_embedding?: boolean;
|
||||
|
|
@ -151,6 +163,7 @@ export interface TrainConfig {
|
|||
differential_guidance_scale?: number;
|
||||
audio_loss_multiplier?: number;
|
||||
max_loss?: number | null;
|
||||
validation_config?: ValidationConfig;
|
||||
}
|
||||
|
||||
export interface QuantizeKwargsConfig {
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
VERSION = "0.10.28"
|
||||
VERSION = "0.10.29"
|
||||
|
|
|
|||
Loading…
Reference in New Issue