File size: 5,352 Bytes
eae424a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | param(
[Parameter(Mandatory = $true)]
[string]$Manifest,
[Parameter(Mandatory = $true)]
[string]$Model,
[Parameter(Mandatory = $true)]
[string]$OutputRoot,
[string[]]$Seeds = @('0', '1', '2'),
[int]$Steps = 12000,
[int]$Images = 2500,
[int]$Layers = 3,
[int]$Size = 224,
[string]$EvaluationSplit = "validation"
)
$ErrorActionPreference = "Stop"
$repo = (Resolve-Path (Join-Path $PSScriptRoot "..")).Path
$manifestPath = (Resolve-Path $Manifest).Path
$modelPath = (Resolve-Path $Model).Path
$manifestHash = (Get-FileHash -LiteralPath $manifestPath -Algorithm SHA256).Hash.ToLowerInvariant()
$modelHash = (Get-FileHash -LiteralPath $modelPath -Algorithm SHA256).Hash.ToLowerInvariant()
$outputPath = [System.IO.Path]::GetFullPath($OutputRoot)
New-Item -ItemType Directory -Force -Path $outputPath | Out-Null
$cargo = (Get-Command cargo).Source
$seedValues = @(
$Seeds | ForEach-Object { $_ -split ',' } | ForEach-Object {
$value = 0L
if (-not [long]::TryParse($_.Trim(), [ref]$value) -or $value -lt 0) {
throw "invalid seed: $_"
}
$value
}
)
if ($seedValues.Count -eq 0 -or @($seedValues | Select-Object -Unique).Count -ne $seedValues.Count) {
throw 'provide at least one seed, without duplicates'
}
function Run-Cargo(
[string[]]$Arguments,
[string]$Stdout,
[string]$Stderr,
[string]$LogLevel
) {
$env:RUST_LOG = $LogLevel
Push-Location $repo
$oldErrorPreference = $ErrorActionPreference
try {
# Windows PowerShell may wrap native stderr as ErrorRecord objects;
# Continue preserves the native exit code while the two streams are
# retained independently.
$ErrorActionPreference = "Continue"
& $cargo @Arguments 1> $Stdout 2> $Stderr
$exitCode = $LASTEXITCODE
} finally {
$ErrorActionPreference = $oldErrorPreference
Pop-Location
}
if ($exitCode -ne 0) {
throw "cargo exited with code $exitCode; see $Stdout and $Stderr"
}
}
foreach ($seed in $seedValues) {
$runDir = Join-Path $outputPath "${Layers}l${Size}-seed${seed}"
New-Item -ItemType Directory -Force -Path $runDir | Out-Null
$decoder = Join-Path $runDir "decoder.bin"
$training = Join-Path $runDir "training.json"
$quality = Join-Path $runDir "quality-imagenette-$EvaluationSplit.json"
if (-not ((Test-Path -LiteralPath $training) -and (Test-Path -LiteralPath $decoder))) {
Write-Output "training seed $seed at $(Get-Date -Format o)"
Run-Cargo `
-Arguments @(
"run", "--release", "--example", "train_decoder", "--",
$manifestPath, $modelPath, "$Steps", "$Images", "$Layers", "$Size", "$seed", $runDir
) `
-Stdout (Join-Path $runDir "training.stdout.log") `
-Stderr (Join-Path $runDir "training.stderr.log") `
-LogLevel "info"
} else {
Write-Output "seed $seed training outputs already exist; preserving them"
}
$record = Get-Content -Raw -LiteralPath $training | ConvertFrom-Json
if (
$record.seed -ne $seed -or
$record.steps -ne $Steps -or
$record.requested_images -ne $Images -or
$record.encoder_layers -ne $Layers -or
$record.image_size -ne $Size -or
$record.dataset_manifest_sha256 -ne $manifestHash -or
$record.model_sha256 -ne $modelHash
) {
throw "existing training record does not match requested cell: $training"
}
if (-not (Test-Path -LiteralPath $quality)) {
Write-Output "evaluating seed $seed at $(Get-Date -Format o)"
Run-Cargo `
-Arguments @(
"run", "--release", "--example", "evaluate_decoder", "--",
"--manifest", $manifestPath,
"--model", $modelPath,
"--decoder", $decoder,
"--split", $EvaluationSplit,
"--layers", "$Layers",
"--size", "$Size",
"--output", $quality,
"--samples", "12"
) `
-Stdout (Join-Path $runDir "evaluation.stdout.log") `
-Stderr (Join-Path $runDir "evaluation.stderr.log") `
-LogLevel "warn"
} else {
Write-Output "seed $seed quality output already exists; preserving it"
}
$qualityRecord = Get-Content -Raw -LiteralPath $quality | ConvertFrom-Json
if (
$qualityRecord.split -ne $EvaluationSplit -or
$qualityRecord.encoder_layers -ne $Layers -or
$qualityRecord.image_size -ne $Size -or
$qualityRecord.manifest_sha256 -ne $manifestHash -or
$qualityRecord.model_sha256 -ne $modelHash -or
$qualityRecord.decoder_sha256 -ne $record.decoder_sha256
) {
throw "quality record does not match its training cell: $quality"
}
$safeTensors = Join-Path $runDir "decoder.safetensors"
if (-not (Test-Path -LiteralPath $safeTensors)) {
& python (Join-Path $repo "tools\convert_decoder.py") `
--training $training `
--input $decoder `
--output $safeTensors
if ($LASTEXITCODE -ne 0) {
throw "decoder SafeTensors conversion failed for seed $seed"
}
}
}
Write-Output "decoder matrix complete at $(Get-Date -Format o)"
|