El desarrollo de este proyecto se ha realizado apoyándose en la implementación original de la arquitectura en PyTorch [ver] y, en mayor medida, en esta adaptación realizada a Tensorflow 1.13 [ver].
La red diseñada ha sido entrenada utilizando los mismos datasets de entrenamiento y validación, DIV2K, que en las dos implementaciones mencionadas. Como datasets de test, se han empleado cuatro conjuntos de imágenes reconocidos en la literatura: Set5, Set14, Urban100 y Manga109, disponibles públicamente aquí. Las imágenes de los datasets se almacenan en formato .bmp para evitar pérdidas por compresión (como las originadas por el formato .jpg).
El entrenamiento se ha realizado de la forma más similar posible al entrenamiento original de la red con el objetivo de replicar sus resultados. Por tanto, los parámetros del entrenamiento son:
- Número de iteraciones (training steps): 1.000.000.
- Optimizador: Adam con parámetros por defecto.
- Learning rate: 1e-4, que cada 200.000 iteraciones se reduce a la mitad.
- Validación: cada 20.000 training steps.
- Data augmentation:
- Volteos aleatorios tanto en el eje X como en el eje Y.
- Rotaciones aleatorias en múltiplos de 90: 0, 90, 180 o 270.
Las métricas empleadas para la medición de la calidad de las reconstrucciones son: PSNR y SSIM, obteniendo unos valores parejos a los otros dos trabajos:
| Arquitectura | Métrica | Set5 | Set14 | Urban100 | Manga109 |
|---|---|---|---|---|---|
| RCAN (Original) | PSNR | 34,74 | 30,65 | 29,09 | 34,44 |
| SSIM | 0,9299 | 0,8482 | 0,8702 | 0,9499 | |
| RCAN (TF 1.13) | PSNR | 34,75 | 30,61 | 29,03 | 34,49 |
| SSIM | 0,9302 | 0,8470 | 0,8693 | 0,9500 | |
| RCAN (TF 2.4 - nuestro) | PSNR | 34,68 | 30,31 | 28,90 | 34,31 |
| SSIM | 0,9300 | 0,8474 | 0,8679 | 0,9494 |
Los módulos necesarios para ejecutar el proyecto son los que se indican a continuación:
- python == 3.7.9
- tensorflow-gpu == 2.4
- open-cv == 4.4.0
- pandas == 1.1.4
- numpy == 1.18.3
- scikit-learn == 0.23.2
- imgaug == 0.4.0
- pillow (PIL) == 8.0.1
El proyecto se encuentra estructurado en las siguientes carpetas:
En este directorio se encuentran todos los datasets empleados en el entrenamiento y testeo de las redes elaboradas, además de otros datasets no utilizados pero también disponibles (General-100, 91-images, BSDS200 y openImages_light).
Cada dataset posee su propio directorio en el cual se encuentran todas sus imágenes.
Todos estos datasets poseen una copia en menor resolución (factor de reducción igual a 3) en la carpeta 'data/lr_images/'. Estas imágenes fueron generadas a través de MATLAB (interpolación bicúbica) y ubicadas en esa dirección de forma manual. Por cada dataset que se almacene, se deberá añadir a 'data/lr_images/' su versión en baja resolución.
En la carpeta 'data/0_csvs/' se podrá encontrar un fichero .csv por cada dataset disponible. Estos ficheros contienen una única columna con identificador 'path' e indica la ruta hacia cada una de las imágenes del dataset, partiendo siempre desde el directorio 'data/'.
Se incluye además un script, a2_create_csv_files.py, que, al ser ejecutado, crea un fichero csv para cada dataset (directorio en la carpeta 'data' que no se llame '0_csvs' o 'lr_images') y lo almacena en 'data/0_csvs/'.
Directorio donde se almacenan los módulos Python implementados para el desarrollo del proyecto. Éstos son:
- constants: en este módulo se definen una serie de valores constantes o identificadores que son posteriormente utilizados por el resto de módulos.
- custom_callbacks: en este módulo se encuentran los Callbacks y funciones que han sido necesarias implementar para realizar el seguimiento de los entrenamientos.
- custom_loss_functions: contiene las implementaciones de las distintas loss functions disponibles para entrenar (mae, mse y sobel_loss). Incluye además una función, 'get_mix_loss_function', que permite combinar de forma ponderada dos loss functions en una sola, y otra función que permite seleccionar una loss function a partir de su identificador en el módulo 'constants'.
- custom_metric_functions: contiene las implementaciones de las distintas loss functions disponibles para entrenar (mae, mse y sobel_loss). Incluye además una función que permite seleccionar una metric function a partir de su identificador en el módulo 'constants'.
- data_processing_functions: contiene las funciones de mapeado necesarias para la carga de los datasets en memoria y su transformaciones.
- PrepareDataset: implementa la clase PrepareDataset, que permite realizar la carga de los datasets de una forma más sencilla al automatizar el proceso mediante parámetros en su constructor. Espera recibir inicialmente los datasets como un único fichero .csv (como los almacenados en 'data/0_csvs/'). Este csv contendrá las rutas hacia las imágenes en alta resolución y, a partir de éstas, y conociendo que las imágenes en baja resolución deben estar almacenadas en 'data/lr_images/', generará las rutas hacia las imágenes de baja resolución.
- psnr_ssim: implementa las métricas PSNR y SSIM correctas para la evaluación. Las proporcionadas por Tensorflow son útiles para entrenar, pero no proporcionan los resultados correctos que permitan comparar el rendimiento de la red con el de otras redes en la literatura.
- RCAN: implementación de la arquitectura RCAN, basada en su mayoría en la implementación hecha para Tensorflow 1.13. Contiene la función 'get_RCAN', que genera una red con el número de grupos y bloques residuales que se especifiquen por parámetros. Además, permite especificar el modo en el que la red se entrenará, es decir, lo que ocurre en cada training step.
- RCAN_inception y RCAN_lkrelu: implementaciones de pequeñas modificaciones sobre la arquitectura original que no mostraron diferencia en los resultados (en pruebas sobre redes de menor tamaño, las RCAN-SCALE3-SHORT).
- training_loops: contiene la definición de los bucles de entrenamiento para las redes RCAN. Estos bucles se definen como clases que heredan de 'keras.Model' y se sobreescriben únicamente los métodos 'train_step' y 'test_step', especificando la forma en la que se desee que se calcule y optimice el valor del loss. Al implementarse de esta manera y no como un bucle a secas, se aprovechan las optimizaciones de código que Tensorflow aplica de forma interna.
En este directorio se almacenan los ficheros de configuración de las redes entrenadas. Éstos son los ficheros que contienen la estructura de la red (.json) y los que contienen los valores de los pesos de las redes (.h5).
Habrá tantas carpetas como estructuras de redes se hayan probado en entrenamientos. La forma de definir si una red pertenece a una estructura u otra se realiza a partir del nombre del script o notebook donde se implementa el entrenamiento. Los nombres de estos archivos seguirán el siguiente formato: identificador numérico, 'train', estructura de la red, características del entrenamiento. Estos cuatro atributos deberán estar separados por '_', siendo siempre los tres primeros elementos el identificador, 'train' y la estructura. En las 'características del entrenamiento' pueden haber más caracteres '_'. Un ejemplo de esta nomenclatura es:
nb05_train_rcan-scale3_loss-computed-on-lr.ipynb
En cada carpeta correspondiente a cada estructura definida, habrá tantos ficheros .json como configuraciones de entrenamiento probadas por cada estructura. Además, por cada entrenamiento habrá una serie de ficheros .h5 que contendrán las configuraciones de pesos que mejores valores obtuvieron en validación para las métricas PSNR y SSIM.
En este directorio se guardan los ficheros .csv generados por el Callback 'Save_Training_Evolution' y contienen la evolución de todas las métricas en cada entrenamiento.
Mediante el script see_metrics_evolution_in_training.py se podrá visualizar gráficamente la evolución de los entrenamientos, además de poder comparar varios entrenamientos entre sí.
En esta carpeta se guardan los scripts que se emplearon originariamente para realizar los primeros entrenamientos en los que la estructura de la red no estaba completamente pulida o los hiper-parámetros del entrenamiento no eran iguales a los de los trabajos mencionados.
Las redes resultantes de estos entrenamientos se pueden ver en el directorio 'TRAINED_MODELS_NEW/', en las estructuras 'RCAN-BASE-SCALE3' y 'RCAN-SCALE3-SHORT'.
Estos scripts no han sido reestructurados (se eliminó y reorganizó gran parte del código) y, por tanto, no pueden ser ejecutados porque darán error de importación o asignación de parámetros.
En este directorio se almacenan las imágenes de los datasets (una carpeta por dataset) reconstruidas (a partir de las imágenes en baja resolución) por la red que se esté evaluando una vez haya sido entrenada.
Este script, mencionado anteriormente, permite visualizar y/o comparar la evolución que muestra un entrenamiento, ya sea mientras éste se ejecuta o una vez finalizado, mediante la consulta del fichero .csv que se va generando conforme el entrenamiento progresa (para que se genere dicho .csv es necesario que en el entrenamiento se incluya el Callback Save_Training_Evolution).
Para más información acerca de los parámetros de este script, consultar su descripción inicial o ejecutar:
python see_metrics_evolution_in_training.py -h
Estos dos archivos contienen la definición del entrenamiendo de la red RCAN propuestas. La diferencia entre ambos se trata de que, al producirse un corte en el suministro eléctrico durante el entrenamiento, se creó un nuevo notebook desde el cual se reanudó el entrenamiento, modificando los parámtros correspondientes a la duración e inicialización del mismo.
En el entrenamiento base, se realizan los siguientes pasos:
- Se construyen unos datasets de entrenamiento y validación empleando los ficheros .csv correspondientes a los datasets Train y Val de DIV2K. Se establece un batch size de 16 y un tamaño de recorte de 48x48 en el dataset de entrenamiento (se aplican recortes a las imágenes con el objetivo de reducir la carga de trabajo de la GPU en el entrenamiento y además, aumentar virtualmente, aparte del data-augmentation, el número de imágenes para entrenar). En el dataset de validación no se aplica ningún recorte y el batch size es igual a 1.
- Se inicializa la red RCAN con el mismo número de grupos y bloques residuales que se indica en la implementación original.
- Se establecen los demás parámetros del entrenamiento:
- Optimizador Adam.
- Número de épocas acorde para que se realicen 1.000.000 de training steps.
- Asignación del learning rate inicial a 1e-4, con una reducción en 1/2 cada 200.000 training steps.
Estos dos notebooks poseen una estructura idéntica a los que se acaban de mencionar, únicamente cambiando la creación/inicialización de la red RCAN. En estos entrenamientos se aplican otros bucles de entrenamiento que hacen que la loss function se pueda calcular sobre sólo la imagen en baja resolución, o sobre imagen de baja y alta resolución simultáneamente.
Este notebook permite evaluar sobre los datasets de test disponibles el rendimiento de la red seleccionada.
La evaluación se realiza con las métricas PSNR y SSIM implementadas en el módulo 'psnr_ssim.py' ya que tienen un funcionamiento más preciso.