Couche SyncBN pour le parallélisme de donnée

Les couches BatchNorm permettent d'accélérer grandement l'apprentissage d'un modèle en lui permettant de converger plus rapidement vers un meilleur optimum. Voir le papier sur la normalisation de couche.

La couche de BatchNormalization applique une transformation qui maintient la moyenne des sorties de la couche précédente à 0 et l'écart type des sorties à 1. Autrement dit, elles calculent des facteurs afin de normaliser les sorties de chaque couche (ou certaines couches) du réseau de neurones. Ces facteurs sont appris lors de l'étape d'apprentissage. À chaque itération de batch, la couche calcule aussi les moyennes et les écarts type du batch (par dimension). La combinaison de tous ces facteurs permet de conserver la moyenne des sorties proche de 0, et l'écart type des sorties proche de 1.

La Batch Normalization travaille différemment pendant l'apprentissage ou pendant la validation ou l'inférence. Il est donc important d'indiquer au modèle son état (apprentissage ou validation).

Pendant l'apprentissage, la couche normalise ses sorties en utilisant la moyenne et l'écart type du batch d'entrée. Plus exactement, la couche retourne (batch - mean(batch)) / (var(batch) + epsilon) * gamma + beta , avec :

• epsilon, une petite constante pour éviter la division par 0,
• gamma, un facteur appris (entraîné) avec un calcul de gradient lors de la backpropagation et qui est initialisé à 1,
• beta, un facteur appris (entraîné) avec un calcul de gradient lors de la backpropagation et qui est initialisé à 0.

Pendant l'inférence ou la validation, la couche normalise ses sorties en utilisant en plus des gamma et beta entraînés, les facteurs moving_mean et moving_var : (batch - moving_mean) / (moving_var + epsilon) * gamma + beta.

moving_mean et moving_var sont des facteurs non entraînés, mais qui sont mis à jour à chaque itération de batch lors de l'apprentissage, selon la méthode suivante :

• moving_mean = moving_mean * momentum + mean(batch) * (1 - momentum)
• moving_var = moving_var * momentum + var(batch) * (1 - momentum)

Lors du parallélisme de données, une réplique du modèle est chargée dans chaque équipement (GPU). Ces répliques sont censées être complètement équivalentes pour chaque équipement. Cependant avec la BatchNormalization, et dû au fait que chaque équipement parallélisé voit passer des mini-batchs différents, il est probable que les facteurs divergent, notamment les variables moving_mean et moving_var.

Si les tailles de mini batch par GPU sont assez grandes, cette divergence peut être jugée acceptable. Cependant il est conseillé et parfois nécessaire de remplacer les couches BatchNorm par des couches SyncBN.

Les couches SyncBN permettent une synchronisation des équipements (lors du parallélisme de données) pour le calcul des facteurs de normalisation.

Une couche SyncBN se définit de la manière suivante dans l'architecture du modèle :

syncBN_layer = torch.nn.SyncBatchNorm(num_features, eps=1e-05, momentum=0.1, affine=True,
                             track_running_stats=True, process_group=None)

Cependant il est possible avec la méthode convert_sync_batchnorm de convertir un modèle existant, en transformant toutes les couches BatchNorm en couches SyncBatchNorm.

Il est nécessaire d'appliquer le convert_sync_batchnorm, avant de transformer le modèle en modèle DDP.

# network is nn.BatchNorm layer
sync_bn_network = torch.nn.SyncBatchNorm.convert_sync_batchnorm(network)
# only single gpu per process is currently supported
ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(
                        sync_bn_network,
                        device_ids=[args.local_rank],
                        output_device=args.local_rank)

Source: Documentation Pytorch

Pour TensorFlow, les couches SyncBN sont encore à un stade expérimental, nous nous permettons de vous renvoyer directement vers les documentations :

• Documentation TensorFlow Keras
• Documentation Horovod