Référence de l'espace de nommage aidge_in_space.training.loop

Boucle d'entraînement de Ballé 2017, portée des cellules 4, 15 et 17 du notebook. Plus de détails...

Classes
class	TrainResult
	Résultat d'un appel à train(). Plus de détails...

Fonctions
None	_say (*args)
	print() immédiat (la sortie est redirigée vers train.log et lue pendant l'entraînement).
None	seed_everything (int seed)
	Fixe la graine de random, numpy et torch (CUDA compris).
torch.device	select_device (str prefer="auto")
	Choisit le device et applique les réglages de performance.
Balle2017	init_reference_weights (Balle2017 model)
	Initialisation du notebook : xavier_normal_ de gain √2 et biais 0,01 sur les conv et deconv.
dict	forward_train (Balle2017 model, torch.Tensor x)
	Passe avant du notebook : bruit uniforme en entraînement, vrai arrondi en évaluation.
	rd_loss (dict out, torch.Tensor x, float lmbda)
	Perte débit-distorsion λ·MSE + bpp, sans synchronisation GPU.
tuple	evaluate (Balle2017 model, list images, device)
	PSNR et débit moyens sur des images, avec vraie quantification.
None	_warn_16bit (int n)
	Avertit que n images à plus de 8 bits ont été mises à l'échelle (rien si n est nul).
torch.Tensor	build_patch_bank (list paths_, int n_patches, int crop, int in_ch, device, seed=None, dict|None stats=None)
	Extrait une fois pour toutes des patchs uint8 (avec le sous-échantillonnage de l'article).
torch.Tensor	sample_batch (torch.Tensor bank, int batch_size)
	Lot tiré avec remise dans la banque, converti en [0, 1], avec retournements aléatoires.
Path	checkpoint_path (TrainConfig cfg, job_id=None, ckpt_dir=None)
	Chemin du checkpoint d'une configuration (motif reconnu par le banc d'essai).
Path	save_checkpoint (Path path, model, optimizer, TrainConfig cfg, dict history, int step, dict run_info)
	Écrit un checkpoint de façon atomique (fichier temporaire puis os.replace).
dict	load_checkpoint (Path path, map_location="cpu")
	Lit un checkpoint avec torch.load(weights_only=True).
str	_read_ref (Path gitdir, str ref)
	Empreinte d'une référence Git (fichier de la référence, puis packed-refs), ou chaîne vide.
str	git_commit ()
	Empreinte (12 caractères) du commit courant, lue sans sous-processus.
None	_warn_callback (int step, Exception exc)
	Avertit qu'un rappel de métriques a échoué (l'entraînement continue).
None	_check_resume (dict ck, TrainConfig cfg, dict run_info, Path ckpt_path)
	Refuse de reprendre un checkpoint qui n'est pas celui de ce run.
tuple	_pick_data (TrainConfig cfg)
	Images d'entraînement et d'évaluation, avec les replis de la conception (§ 2.3.1).
TrainResult	train (TrainConfig cfg, Path ckpt_path, *, dict|None run_info=None, device=None, Callable[[], bool]|None should_stop=None, Callable[[int, dict], None]|None on_metrics=None)
	Entraîne (ou reprend) un modèle Ballé 2017 et sauvegarde ses checkpoints.

Variables
int	LIKELIHOOD_MIN = 1e-9
	Borne inférieure des probabilités (comme le notebook) avant -log2.
tuple	HISTORY_KEYS = ("step", "loss", "bpp", "psnr", "eval_step", "eval_psnr", "eval_bpp")
	Clés de l'historique, celles du notebook.
	REPLACE_RETRIES
	Nouvelles tentatives de os.replace sous Windows (fichier ouvert en lecture) et délai entre deux, en secondes.
	REPLACE_DELAY_S
	Nouvelles tentatives de os.replace sous Windows (fichier ouvert en lecture) et délai entre deux, en secondes.
tuple	ARCH_KEYS = ("in_ch", "N", "gdn_last", "lmbda")
	Les paramètres de l'architecture doivent être identiques à ceux du checkpoint repris.

Description détaillée

Boucle d'entraînement de Ballé 2017, portée des cellules 4, 15 et 17 du notebook.

Banque de patchs uint8 pré-extraits sur le device, perte λ·MSE + bpp, Adam avec pas divisé par 10 en fin d'entraînement, évaluation avec vraie quantification, checkpoints atomiques et reprise. Le contrat est dans docs/conception_entrainement.md (§ 2.3).

Documentation des fonctions

_check_resume()

None aidge_in_space.training.loop._check_resume ( dict ck, TrainConfig cfg, dict run_info, Path ckpt_path )

protected

Refuse de reprendre un checkpoint qui n'est pas celui de ce run.

Exceptions

ConfigError

si la clé run manque, si le job ou le nom diffère, ou si l'architecture ou λ diffère.

_pick_data()

tuple aidge_in_space.training.loop._pick_data ( TrainConfig cfg )

protected

Images d'entraînement et d'évaluation, avec les replis de la conception (§ 2.3.1).

Renvoie

Triplet (chemins d'entraînement, chemins d'évaluation, source de l'évaluation).

Exceptions

ConfigError

si le dossier d'entraînement désigné est vide, ou si les deux jeux sont vides.

build_patch_bank()

torch.Tensor aidge_in_space.training.loop.build_patch_bank	(	list	paths_,
		int	n_patches,
		int	crop,
		int	in_ch,
			device,
			seed = None,
		dict | None	stats = None )

Extrait une fois pour toutes des patchs uint8 (avec le sous-échantillonnage de l'article).

Par image : un facteur d'échelle aléatoire entre max(crop / min(w, h), 0,3) et 0,75, puis plusieurs recadrages. Si le jeu compte plus d'images que de patchs, n_patches chemins sont tirés au hasard (reproductible avec seed) avant tout décodage : la banque ne dépasse jamais n_patches patchs. Les images à plus de 8 bits sont mises à l'échelle (to_8bit).

Paramètres

paths_	Chemins des images.
n_patches	Nombre de patchs visé (borne haute de la banque).
crop	Côté des patchs.
in_ch	1 (gris) ou 3 (RGB).
device	Device de la banque.
seed	Graine du tirage des chemins (défaut : générateur global de random).
stats	Si fourni, reçoit le nombre d'images converties dans stats["converted"] et l'avertissement 16 bits est laissé à l'appelant ; sinon il est écrit ici.

Renvoie

Tenseur uint8 (n, C, crop, crop) sur device.

Exceptions

ConfigError

si aucune image n'est fournie ou lisible.

checkpoint_path()

Path aidge_in_space.training.loop.checkpoint_path	(	TrainConfig	cfg,
			job_id = None,
			ckpt_dir = None )

Chemin du checkpoint d'une configuration (motif reconnu par le banc d'essai).

Paramètres

cfg	Configuration.
job_id	Numéro du job (suffixe _run{N}), ou None (nom du notebook).
ckpt_dir	Dossier (défaut config.CKPT_DIR, lu à l'appel).

Renvoie

Chemin bls2017_c{in_ch}_N{N}_lambda{λ}[_run{job}].pt.

evaluate()

tuple aidge_in_space.training.loop.evaluate	(	Balle2017	model,
		list	images,
			device )

PSNR et débit moyens sur des images, avec vraie quantification.

Chaque image est complétée par réplication à un multiple de 16 ; l'image reconstruite est rognée, écrêtée puis ré-arrondie sur 8 bits avant le PSNR. Le modèle reste en mode évaluation.

Paramètres

model	Modèle.
images	Tenseurs 1 x C x H x W dans [0, 1].
device	Device du modèle.

Renvoie

Couple (PSNR moyen en dB, débit moyen en bpp).

forward_train()

dict aidge_in_space.training.loop.forward_train	(	Balle2017	model,
		torch.Tensor	x )

Passe avant du notebook : bruit uniforme en entraînement, vrai arrondi en évaluation.

Paramètres

model	Modèle (son mode train/eval choisit la quantification).
x	Lot B x C x H x W dans [0, 1].

Renvoie

Dictionnaire y, y_hat, x_hat et likelihoods (probabilités non bornées).

git_commit()

str aidge_in_space.training.loop.git_commit

(

)

Empreinte (12 caractères) du commit courant, lue sans sous-processus.

Gère le dossier .git, le fichier .git d'un worktree, HEAD détaché et packed-refs.

Renvoie

12 caractères hexadécimaux, ou chaîne vide en cas d'échec (l'image Docker n'a pas git).

init_reference_weights()

Balle2017 aidge_in_space.training.loop.init_reference_weights

(

Balle2017

model

)

Initialisation du notebook : xavier_normal_ de gain √2 et biais 0,01 sur les conv et deconv.

La GDN et le prior gardent leur initialisation.

Paramètres

model

Modèle, modifié en place.

Renvoie

Le même modèle.

load_checkpoint()

dict aidge_in_space.training.loop.load_checkpoint	(	Path	path,
			map_location = "cpu" )

Lit un checkpoint avec torch.load(weights_only=True).

Paramètres

path	Fichier.
map_location	Device de destination des tenseurs.

Renvoie

Le dictionnaire du checkpoint.

Exceptions

FileNotFoundError

si le fichier est absent.

rd_loss()

aidge_in_space.training.loop.rd_loss	(	dict	out,
		torch.Tensor	x,
		float	lmbda )

Perte débit-distorsion λ·MSE + bpp, sans synchronisation GPU.

Paramètres

out	Sortie de forward_train().
x	Lot d'entrée.
lmbda	Poids de la distorsion.

Renvoie

Triplet de scalaires (perte, MSE, bpp).

sample_batch()

torch.Tensor aidge_in_space.training.loop.sample_batch	(	torch.Tensor	bank,
		int	batch_size )

Lot tiré avec remise dans la banque, converti en [0, 1], avec retournements aléatoires.

Paramètres

bank	Banque uint8 de build_patch_bank().
batch_size	Taille du lot.

Renvoie

Tenseur float sur le device de la banque.

save_checkpoint()

Path aidge_in_space.training.loop.save_checkpoint	(	Path	path,
			model,
			optimizer,
		TrainConfig	cfg,
		dict	history,
		int	step,
		dict	run_info )

Écrit un checkpoint de façon atomique (fichier temporaire puis os.replace).

Le contenu est relisible avec torch.load(weights_only=True). Sous Windows, os.replace est retenté si le fichier est ouvert en lecture (le banc d'essai).

Paramètres

path	Fichier cible.
model	Modèle.
optimizer	Optimiseur.
cfg	Configuration.
history	Historique (clés du notebook).
step	Pas atteint.
run_info	Clé run (sans date, ajoutée ici).

Renvoie

path.

Exceptions

PermissionError

si le remplacement échoue après toutes les tentatives.

seed_everything()

None aidge_in_space.training.loop.seed_everything

(

int

seed

)

Fixe la graine de random, numpy et torch (CUDA compris).

Paramètres

seed

Graine.

select_device()

torch.device aidge_in_space.training.loop.select_device

(

str

prefer = "auto"

)

Choisit le device et applique les réglages de performance.

Sur CUDA : cudnn.benchmark et TF32. Sur CPU : tous les cœurs sauf un (pour laisser l'interface réactive).

Paramètres

prefer

"auto" (CUDA si disponible), "cpu" ou "cuda".

Renvoie

Le device.

Exceptions

ConfigError

si CUDA est demandé mais absent, ou si prefer est inconnu.

train()

TrainResult aidge_in_space.training.loop.train	(	TrainConfig	cfg,
		Path	ckpt_path,
		*	,
		dict | None	run_info = None,
			device = None,
		Callable[[], bool] | None	should_stop = None,
		Callable[[int, dict], None] | None	on_metrics = None )

Entraîne (ou reprend) un modèle Ballé 2017 et sauvegarde ses checkpoints.

Déroulement (docs/conception_entrainement.md, § 2.3.1) : données, reprise éventuelle, boucle Adam, puis journal, évaluation et sauvegarde périodiques. should_stop est appelé à la fin de chaque pas ; s'il renvoie vrai, le checkpoint est sauvegardé au pas courant.

Paramètres

cfg	Configuration.
ckpt_path	Checkpoint à écrire (et à reprendre s'il existe).
run_info	Clé run du checkpoint sans date (défaut : job None, nom et auteur de cfg, commit courant).
device	Device (défaut : select_device("auto")).
should_stop	Appelable sans argument, vrai pour demander l'arrêt.
on_metrics	Appelable (pas, dictionnaire kind=train ou kind=eval).

Renvoie

Le résultat, de statut "done" ou "cancelled".

Exceptions

ConfigError

si les données manquent ou si le checkpoint existant est incompatible.