Выложил https://github.com/DeliriumV01D/TransGANpp игрушечный пример обучения целиком трансформерной GAN на данных MNIST dataset. Имеется train loop для генератора и дискриминатора. Именование модулей чтобы потом при инициализации весов различать их по типу. Добавлена OpenCV визуализация выхода сети - для этого делается конвертация torch::Tensor -> cv::Mat. Подсчет количества весов (чтоб было меньше чем для DCGAN выполняющего эту же задачу).
Используется классический необучаемый позиционный энкодер собственной реализации. Инициализация весов трансформеров в данную секунду осуществляется методом xavier_normal. Для стабилизации обучения в случае если генератор побеждает, то пропускаем апдейт генератора, пока дискриминатор не вернет лидерство. Имеется управление скоростью обучения. Переход от изображения к последовательности токенов осуществляется посредством n сверток 16 на 16 со страйдом на размер свертки и паддингом 6. Из изображения 28 на 28 получается n изображений 2 на 2. Из которых, в свою очередь, собирается 4 токена длины n. Число блоков трансформеров взято равное 3. Количество голов механизма внимания - 4. Метод оптимизации - adam(0.5, 0,9).