{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.18","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":13836,"databundleVersionId":1718836,"sourceType":"competition"}],"dockerImageVersionId":31091,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#Importamos las librerias necesarias\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as transforms\nimport timm\nimport gc\nimport os\nimport time\nimport random\nfrom datetime import datetime\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nfrom sklearn import model_selection, metrics\nimport os\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nimport seaborn as sns","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-07-22T13:53:51.918576Z","iopub.execute_input":"2025-07-22T13:53:51.918905Z","iopub.status.idle":"2025-07-22T13:54:29.535184Z","shell.execute_reply.started":"2025-07-22T13:53:51.918878Z","shell.execute_reply":"2025-07-22T13:54:29.529239Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Tomamos los datos de Kaggle y los cargamos en un DataFrame\n\nData = \"/kaggle/input/cassava-leaf-disease-classification\"\nTrain = \"/kaggle/input/cassava-leaf-disease-classification/train_images/\"\nTest = \"/kaggle/input/cassava-leaf-disease-classification/test_images/\"\n\n\ndf = pd.read_csv(os.path.join(Data, \"train.csv\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T13:54:57.049986Z","iopub.execute_input":"2025-07-22T13:54:57.050453Z","iopub.status.idle":"2025-07-22T13:54:57.085077Z","shell.execute_reply.started":"2025-07-22T13:54:57.050425Z","shell.execute_reply":"2025-07-22T13:54:57.079253Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Para este proyecto se hará uso del TPU proporcionado por Kaggle, para lo cual implementaremos primero paralelización del TPU. Esto nos permitirá acelerar el entrenamiento de modelos de deep learning utilizando la librería Torch XLA, que facilita la integración de PyTorch con dispositivos TPU. La paralelización se logra distribuyendo los datos y el proceso de entrenamiento entre los diferentes núcleos del TPU, optimizando así el uso de recursos y reduciendo significativamente el tiempo de entrenamiento. Además, se emplearán técnicas como el uso de DataLoader y ParallelLoader para manejar eficientemente los datos durante el entrenamiento distribuido.","metadata":{}},{"cell_type":"markdown","source":"`XLA_USE_BF16`: Activa el uso de precisión mixta (BF16), lo que acelera el entrenamiento y reduce el consumo de memoria en TPU.\n`XLA_TENSOR_ALLOCATOR_MAXSIZE`: Limita el tamaño máximo del asignador de tensores en XLA, ayudando a controlar el uso de memoria.","metadata":{}},{"cell_type":"code","source":"os.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"\n\n# Para la reproducibilidad de los resultados hacemos uso de una semilla fija\n# Esto asegura que los resultados sean consistentes en diferentes ejecuciones\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n\nseed_everything(42)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:13.918748Z","iopub.execute_input":"2025-07-22T10:58:13.918966Z","iopub.status.idle":"2025-07-22T10:58:13.934708Z","shell.execute_reply.started":"2025-07-22T10:58:13.918946Z","shell.execute_reply":"2025-07-22T10:58:13.928054Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Definición de Parámetros: `IMG_SIZE`, `BATCH_SIZE`, `LR`, y `N_EPOCHS`\n\nEn este apartado, vamos a definir los parámetros clave para configurar el entrenamiento de nuestro modelo Vision Transformer (ViT) para la clasificación de enfermedades en hojas. Estos parámetros influirán directamente en el comportamiento y rendimiento del modelo durante el entrenamiento.\n\n1. **`IMG_SIZE = 224`**:\n   El parámetro `IMG_SIZE` define el tamaño de las imágenes de entrada que se alimentan al modelo. Dado que estamos utilizando un **Vision Transformer (ViT)**, que requiere imágenes de un tamaño específico, hemos decidido establecer un tamaño de imagen de **224x224 píxeles**. Este es un tamaño comúnmente utilizado en muchos modelos preentrenados, lo que ayuda a optimizar el tiempo de entrenamiento sin perder demasiada resolución en las imágenes.\n\n   - **Razón para elegir `224x224`**: Este tamaño es un compromiso entre calidad y eficiencia computacional. Permite capturar suficiente detalle en las imágenes para tareas de clasificación, mientras mantiene un tamaño manejable para los recursos computacionales.\n\n2. **`BATCH_SIZE = 16`**:\n   El `BATCH_SIZE` determina cuántas imágenes serán procesadas en cada iteración del entrenamiento. Con un `batch_size` de **16**, el modelo procesará 16 imágenes simultáneamente en cada paso, actualizando los parámetros basados en el cálculo de los gradientes promedio. Un tamaño de lote más pequeño (como 16) ayuda a reducir el consumo de memoria y puede mejorar la generalización, pero también puede hacer que el entrenamiento sea más ruidoso debido a la mayor variabilidad en los gradientes.\n\n   - **Consideraciones para elegir `16`**: Un tamaño de 16 es una opción común que ofrece un buen equilibrio entre eficiencia computacional y generalización en modelos de visión por computadora. En algunos casos, podría ser necesario ajustar este valor dependiendo del tamaño de las imágenes o la memoria disponible.\n\n3. **`LR = 2e-05`** (Learning Rate):\n   La **tasa de aprendizaje** (`LR`) controla qué tan grande es el paso que da el modelo en cada actualización de los parámetros durante el entrenamiento. Un valor de **2e-05** indica que el modelo hará pequeñas actualizaciones en cada iteración. Esto puede ser útil para evitar que el modelo dé pasos demasiado grandes, lo que podría llevar a una mala convergencia o incluso hacer que el entrenamiento se vuelva inestable.\n\n   - **Razón para elegir `2e-05`**: Este valor es un buen punto de partida basado en modelos preentrenados de ViT. Sin embargo, la tasa de aprendizaje podría ajustarse según el comportamiento observado en el entrenamiento. Si el modelo se entrena demasiado lento, se podría intentar aumentar la tasa de aprendizaje; si el modelo se vuelve inestable, se debería reducir.\n\n4. **`N_EPOCHS = 10`**:\n   El número de **`epochs`** define cuántas veces todo el conjunto de datos se pasará a través del modelo. Un valor de **10** implica que el modelo entrenará durante 10 pasadas completas sobre el conjunto de entrenamiento. El número de epochs es crucial porque un valor bajo podría resultar en un entrenamiento insuficiente, mientras que uno alto podría dar lugar a sobreajuste (overfitting).\n\n   - **Consideraciones para elegir `10`**: En este caso, 10 épocas es un valor inicial razonable para observar cómo el modelo se ajusta a los datos. Dependiendo de los resultados, podríamos aumentar o disminuir este valor. Es importante observar el rendimiento en el conjunto de validación para determinar si el modelo está aprendiendo adecuadamente o si es necesario más entrenamiento.\n\n\n","metadata":{}},{"cell_type":"code","source":"IMG_SIZE = 224\n\nBATCH_SIZE = 16\n\nLR = 2e-05\n\nN_EPOCHS = 10","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:13.935356Z","iopub.execute_input":"2025-07-22T10:58:13.935599Z","iopub.status.idle":"2025-07-22T10:58:13.943724Z","shell.execute_reply.started":"2025-07-22T10:58:13.935576Z","shell.execute_reply":"2025-07-22T10:58:13.940326Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Veamos como son los datos \ndf = pd.read_csv(os.path.join(Data, \"train.csv\"))\n\n# Realizamos la separacion de train y test:\ntrain_df, valid_df = model_selection.train_test_split(df, \n                                                      test_size=0.1, \n                                                      random_state=42, \n                                                      stratify=df.label.values)\n\ndf.head()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:13.945510Z","iopub.execute_input":"2025-07-22T10:58:13.945793Z","iopub.status.idle":"2025-07-22T10:58:13.995025Z","shell.execute_reply.started":"2025-07-22T10:58:13.945770Z","shell.execute_reply":"2025-07-22T10:58:13.991144Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Graficamos el diagrama de barras\ndf.label.value_counts().plot(kind=\"bar\")\n\nlabel_counts = df.label.value_counts()\n\n# Graficamos el pie chart\nplt.figure(figsize=(8, 8))\nplt.pie(label_counts.values, labels=label_counts.index, autopct='%1.1f%%', colors=sns.color_palette(\"Set2\", len(label_counts)))\n\nplt.title('Distribución de Etiquetas en el Conjunto de Datos', fontsize=16)\n\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:13.997175Z","iopub.execute_input":"2025-07-22T10:58:13.997391Z","iopub.status.idle":"2025-07-22T10:58:14.287851Z","shell.execute_reply.started":"2025-07-22T10:58:13.997371Z","shell.execute_reply":"2025-07-22T10:58:14.283514Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Conjunto de Datos Cassava\n\nEl código a continuación define una clase llamada `CassavaDataset`, que facilita la creación de un conjunto de datos personalizado en PyTorch. Esta clase está diseñada específicamente para cargar y preprocesar imágenes del conjunto de datos **Cassava**, que se utilizan para entrenar o evaluar modelos de clasificación de enfermedades en plantas de cassava.\n\n1. **Inicialización (`__init__`)**:\n   - El constructor de la clase recibe varios parámetros: \n     - `df`: un DataFrame que contiene los nombres de las imágenes y sus etiquetas.\n     - `data_path`: la ruta donde se almacenan las imágenes. Por defecto, se utiliza `Data`.\n     - `mode`: un valor que indica si el conjunto de datos es de entrenamiento (\"train\") o prueba (\"test\").\n     - `transforms`: un parámetro opcional que permite aplicar transformaciones de preprocesamiento (como normalización o aumentos de datos) a las imágenes antes de ser procesadas por el modelo.\n   - Dependiendo del valor de `mode`, se selecciona el directorio adecuado (`train_images` o `test_images`) donde se encuentran las imágenes.\n\n2. **Tamaño del Conjunto de Datos (`__len__`)**:\n   - El método `__len__` devuelve el número total de imágenes en el conjunto de datos, que se obtiene del DataFrame `df`. Esto es útil para saber cuántas muestras hay disponibles para el entrenamiento o la evaluación.\n\n3. **Obtención de un Ítem del Conjunto de Datos (`__getitem__`)**:\n   - El método `__getitem__` toma un índice y retorna una imagen y su etiqueta correspondiente.\n   - Primero, se construye la ruta completa de la imagen basándose en el índice y el directorio de imágenes.\n   - Luego, se carga la imagen desde el disco y se convierte a formato RGB.\n   - Si se han definido transformaciones (como redimensionamiento o normalización), estas se aplican a la imagen antes de devolverla.\n   - Finalmente, el método retorna tanto la imagen procesada como su etiqueta.\n\nEn resumen, esta clase permite cargar y transformar eficientemente las imágenes del conjunto de datos Cassava, facilitando su uso en el entrenamiento y evaluación de modelos de aprendizaje automático en PyTorch.\n","metadata":{}},{"cell_type":"code","source":"class CassavaDataset(torch.utils.data.Dataset):\n\n    def __init__(self, df, data_path= Data, mode=\"train\", transforms=None):\n        super().__init__()\n        self.df_data = df.values\n        self.data_path = data_path\n        self.transforms = transforms\n        self.mode = mode\n        self.data_dir = \"train_images\" if mode == \"train\" else \"test_images\"\n\n    def __len__(self):\n        return len(self.df_data)\n\n    def __getitem__(self, index):\n        img_name, label = self.df_data[index]\n        img_path = os.path.join(self.data_path, self.data_dir, img_name)\n        img = Image.open(img_path).convert(\"RGB\")\n\n        if self.transforms is not None:\n            image = self.transforms(img)\n\n        return image, label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:14.289558Z","iopub.execute_input":"2025-07-22T10:58:14.290010Z","iopub.status.idle":"2025-07-22T10:58:14.300418Z","shell.execute_reply.started":"2025-07-22T10:58:14.289987Z","shell.execute_reply":"2025-07-22T10:58:14.296515Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Data Augmentation\n\nEl código a continuación define dos conjuntos de transformaciones para los datos de entrenamiento y validación. Estas transformaciones son aplicadas a las imágenes antes de ser alimentadas al modelo, con el fin de mejorar su generalización y ayudar a prevenir el sobreajuste (overfitting) al introducir variaciones en los datos.\n\n1. **Transformaciones para el Conjunto de Entrenamiento (`transforms_train`)**:\n   - **`Resize`**: Redimensiona las imágenes a un tamaño específico, definido por `IMG_SIZE` (que es un valor previamente definido, como 224x224 píxeles). Esto asegura que todas las imágenes tengan un tamaño uniforme antes de ser procesadas.\n   - **`RandomHorizontalFlip`**: Realiza un volteo horizontal de las imágenes con una probabilidad del 30%. Esto ayuda a que el modelo sea menos sensible a la orientación de las imágenes.\n   - **`RandomVerticalFlip`**: Realiza un volteo vertical de las imágenes con una probabilidad del 30%. De nuevo, esto aumenta la variabilidad de las imágenes y mejora la robustez del modelo.\n   - **`RandomRotation`**: Rota las imágenes aleatoriamente en un rango de hasta 10 grados. Esta transformación es útil para simular variaciones en la orientación de las hojas.\n   - **`RandomAffine`**: Aplica una transformación afín aleatoria en las imágenes, con un rango de 10 grados. Esto puede incluir cambios en la escala y rotación de la imagen.\n   - **`RandomResizedCrop`**: Recorta aleatoriamente una porción de la imagen y la redimensiona a `IMG_SIZE`. Esto ayuda a simular variaciones en el encuadre y la escala de los objetos.\n   - **`ToTensor`**: Convierte la imagen en un tensor de PyTorch, lo que la prepara para ser utilizada en un modelo de aprendizaje automático.\n   - **`Normalize`**: Normaliza las imágenes utilizando los valores de media y desviación estándar de los canales de color (Rojo, Verde y Azul). Esto es importante para que el modelo pueda aprender de manera más eficiente.\n\n2. **Transformaciones para el Conjunto de Validación (`transforms_valid`)**:\n   - **`Resize`**: Similar al conjunto de entrenamiento, redimensiona las imágenes a `IMG_SIZE`.\n   - **`ToTensor`**: Convierte las imágenes en tensores de PyTorch.\n   - **`Normalize`**: Normaliza las imágenes utilizando los mismos valores de media y desviación estándar que en el conjunto de entrenamiento, para asegurar que las imágenes de validación tengan la misma escala que las de entrenamiento.\n","metadata":{}},{"cell_type":"code","source":"# Tiempo para realizar aumentos de datos (Data Augmentation)\n\n# Conjunto de transformaciones para el conjunto de entrenamiento\ntransforms_train = transforms.Compose(\n[\n    # Redimensiona la imagen al tamaño previamente especificado en IMG_SIZE\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    \n    # Voltea aleatoriamente las imágenes de manera horizontal con una probabilidad del 30%\n    transforms.RandomHorizontalFlip(p=0.3),\n    \n    # Voltea aleatoriamente las imágenes de manera vertical con una probabilidad del 30%\n    transforms.RandomVerticalFlip(p=0.3),\n    \n    # Rota aleatoriamente las imágenes hasta 10 grados\n    transforms.RandomRotation(10),\n    \n    # Aplica una transformación afín aleatoria de hasta 10 grados\n    transforms.RandomAffine(10),\n    \n    # Recorta aleatoriamente las imágenes y las redimensiona al tamaño especificado por IMG_SIZE\n    transforms.RandomResizedCrop(IMG_SIZE),\n    \n    # Convierte la imagen en un tensor, adecuado para el modelo en PyTorch\n    transforms.ToTensor(),\n    \n    # Normaliza la imagen usando la media y desviación estándar de los canales RGB\n    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n]\n)\n\n# Conjunto de transformaciones para el conjunto de validación\ntransforms_valid = transforms.Compose(\n[\n    # Redimensiona la imagen al tamaño especificado en IMG_SIZE\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    \n    # Convierte la imagen en un tensor\n    transforms.ToTensor(),\n    \n    # Normaliza la imagen utilizando la media y desviación estándar de los canales RGB\n    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n]\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:14.302451Z","iopub.execute_input":"2025-07-22T10:58:14.302780Z","iopub.status.idle":"2025-07-22T10:58:14.315971Z","shell.execute_reply.started":"2025-07-22T10:58:14.302737Z","shell.execute_reply":"2025-07-22T10:58:14.311247Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Modelo ViTBase16\n\nEl código a continuación define una clase llamada `ViTBase16`, que implementa un modelo **Vision Transformer (ViT)** para la clasificación de imágenes utilizando PyTorch. Este modelo está basado en la arquitectura **ViT Base 16**, con imágenes de entrada redimensionadas a 224x224 píxeles y una división en 16 parches.\n\n1. **Inicialización (`__init__`)**:\n   - El constructor de la clase define el modelo Vision Transformer utilizando la función `timm.create_model` de la biblioteca `timm`, que carga el modelo ViT con parches de tamaño 16x16.\n   - Si se indica `pretrained=True`, se cargan los pesos preentrenados desde una ruta específica, lo que permite transferir el conocimiento de modelos previamente entrenados en grandes conjuntos de datos.\n   - Se ajusta la capa de salida (`model.head`) para que tenga el número de clases especificado por `n_classes`.\n\n2. **Paso hacia adelante (`forward`)**:\n   - El método `forward` define cómo los datos de entrada son pasados a través del modelo para obtener las predicciones. Este método es invocado automáticamente durante el entrenamiento y la validación.\n\n3. **Entrenamiento por Época (`train_one_epoch`)**:\n   - Este método ejecuta una iteración completa de entrenamiento:\n     - Inicializa las variables de pérdida y precisión para la época.\n     - Realiza la pasada hacia adelante para obtener las predicciones.\n     - Calcula la pérdida y la precisión.\n     - Realiza la pasada hacia atrás para actualizar los gradientes.\n     - Ajusta los parámetros del modelo utilizando el optimizador.\n   - Si el dispositivo es de tipo `xla`, se emplea un método optimizado para trabajar en entornos como Google Cloud.\n\n4. **Validación por Época (`validate_one_epoch`)**:\n   - Este método realiza la validación del modelo en una época, donde se calcula la pérdida y precisión sin actualizar los gradientes, permitiendo evaluar el rendimiento del modelo en datos no vistos durante el entrenamiento.\n   - La validación se ejecuta en modo de evaluación para evitar cambios en los pesos del modelo.\n","metadata":{}},{"cell_type":"code","source":"class ViTBase16(nn.Module):\n    def __init__(self, n_classes, pretrained=False, model_path=None):\n        super(ViTBase16, self).__init__()\n\n        # Crear el modelo Vision Transformer\n        self.model = timm.create_model(\"vit_base_patch16_224\", pretrained=False)\n\n        # Cargar pesos preentrenados si se solicita\n        if pretrained:\n            if model_path is not None:\n                self.model.load_state_dict(torch.load(model_path))\n            else:\n                raise ValueError(\"Se requiere un modelo preentrenado y una ruta de archivo si se activa `pretrained`.\")\n\n        # Ajustar la última capa para que tenga el número correcto de clases\n        self.model.head = nn.Linear(self.model.head.in_features, n_classes)\n\n    def forward(self, x):\n        # Pasar las imágenes por el modelo para obtener las predicciones\n        return self.model(x)\n\n    def train_one_epoch(self, train_loader, criterion, optimizer, device, lr_scheduler=None):\n        # Inicializar variables de pérdida y precisión\n        epoch_loss = 0.0\n        epoch_accuracy = 0.0\n\n        # Establecer el modelo en modo de entrenamiento\n        self.model.train()\n\n        for i, (data, target) in enumerate(train_loader):\n            # Enviar los datos al dispositivo adecuado (CPU o GPU)\n            data, target = data.to(device), target.to(device)\n\n            # Limpiar los gradientes del optimizador\n            optimizer.zero_grad()\n\n            # Paso hacia adelante: calcular las predicciones del modelo\n            output = self.forward(data)\n\n            # Calcular la pérdida\n            loss = criterion(output, target)\n\n            # Paso hacia atrás: calcular los gradientes de la pérdida con respecto a los parámetros del modelo\n            loss.backward()\n\n            # Calcular la precisión\n            accuracy = (output.argmax(dim=1) == target).float().mean()\n\n            # Acumular la pérdida y precisión de la época\n            epoch_loss += loss.item()\n            epoch_accuracy += accuracy.item()\n\n            # Realizar un paso de optimización\n            optimizer.step()\n\n            # Actualizar la tasa de aprendizaje si se usa un scheduler\n            if lr_scheduler:\n                lr_scheduler.step()\n\n            # Imprimir información sobre el progreso\n            if i % 20 == 0:\n                print(f\"\\tBATCH {i+1}/{len(train_loader)} - LOSS: {loss.item():.4f} - ACC: {accuracy.item():.4f}\")\n\n        # Promediar la pérdida y precisión para la época\n        return epoch_loss / len(train_loader), epoch_accuracy / len(train_loader)\n\n    def validate_one_epoch(self, valid_loader, criterion, device):\n        # Inicializar variables para la pérdida y precisión de validación\n        valid_loss = 0.0\n        valid_accuracy = 0.0\n\n        # Establecer el modelo en modo de evaluación (sin gradientes)\n        self.model.eval()\n\n        with torch.no_grad():  # Desactivar el cálculo de gradientes para la validación\n            for data, target in valid_loader:\n                # Enviar los datos al dispositivo adecuado\n                data, target = data.to(device), target.to(device)\n    \n                # Paso hacia adelante: calcular las predicciones\n                output = self.model(data)\n    \n                # Calcular la pérdida\n                loss = criterion(output, target)\n    \n                # Calcular la precisión\n                accuracy = (output.argmax(dim=1) == target).float().mean()\n    \n                # Acumular la pérdida y precisión de validación\n                valid_loss += loss.item()\n                valid_accuracy += accuracy.item()\n\n        # Promediar la pérdida y precisión para la validación\n        return valid_loss / len(valid_loader), valid_accuracy / len(valid_loader)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:14.318095Z","iopub.execute_input":"2025-07-22T10:58:14.318310Z","iopub.status.idle":"2025-07-22T10:58:14.335801Z","shell.execute_reply.started":"2025-07-22T10:58:14.318289Z","shell.execute_reply":"2025-07-22T10:58:14.330323Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### `fit_tpu`\n\nEl código a continuación define una función llamada `fit_tpu`, que entrena y valida un modelo en una arquitectura distribuida utilizando TPUs. Esta función maneja tanto el proceso de entrenamiento como el de validación, e implementa técnicas de optimización para mejorar la eficiencia durante el entrenamiento en entornos distribuidos.\n\n1. **Inicialización de Variables**:\n   - Se inicializan las listas `train_losses`, `valid_losses`, `train_accs`, y `valid_accs` para almacenar las métricas de pérdida y precisión durante cada época.\n   - Se establece la variable `valid_loss_min` como `np.inf`, que se usará para hacer un seguimiento de la menor pérdida de validación observada.\n\n2. **Ciclo de Entrenamiento**:\n   - La función realiza un ciclo de entrenamiento que se repite por el número de épocas especificado (`epochs`).\n   - En cada época, se realiza lo siguiente:\n     - **Entrenamiento**: Utiliza el método `train_one_epoch` para entrenar el modelo en los datos de entrenamiento (`train_loader`). Durante el entrenamiento, se calcula la pérdida y la precisión, y se imprime esta información para cada época.\n     - **Validación**: Si se proporciona un conjunto de datos de validación (`valid_loader`), se realiza la validación utilizando el método `validate_one_epoch`. Durante la validación, se calcula la pérdida y la precisión, y se imprime esta información.\n   \n3. **Manejo de Memoria**:\n   - Se usa `gc.collect()` después de cada operación importante para liberar la memoria no utilizada y optimizar el uso de recursos, especialmente cuando se entrena en entornos distribuidos como TPUs.\n\n4. **Guardado del Mejor Modelo**:\n   - Si la pérdida de validación mejora (es decir, es menor que la pérdida mínima registrada), se guarda el modelo actual. Esto permite mantener la mejor versión del modelo durante el entrenamiento.\n\n5. **Devolución de Resultados**:\n   - Al final de las épocas, la función devuelve un diccionario que contiene las pérdidas y precisiones de entrenamiento y validación para cada época. Esto permite hacer un seguimiento detallado del rendimiento del modelo a lo largo del tiempo.\n\n","metadata":{}},{"cell_type":"code","source":"def fit_tpu(model, epochs, device, criterion, optimizer, train_loader, valid_loader=None):\n    import torch_xla.core.xla_model as xm\n    import torch_xla.distributed.parallel_loader as pl\n    \n    valid_loss_min = np.inf\n    train_losses, valid_losses = [], []\n    train_accs, valid_accs = [], []\n\n    for epoch in range(1, epochs + 1):\n        gc.collect()\n        para_train_loader = pl.ParallelLoader(train_loader, [device])\n\n        xm.master_print(f\"{'='*50}\")\n        xm.master_print(f\"EPOCH {epoch} - TRAINING...\")\n\n        train_loss, train_acc = model.train_one_epoch(\n            para_train_loader.per_device_loader(device), criterion, optimizer, device\n        )\n        \n        xm.master_print(f\"\\n\\t[TRAIN] EPOCH {epoch} - LOSS: {train_loss:.4f}, ACCURACY: {train_acc:.4f}\\n\")\n        train_losses.append(train_loss)\n        train_accs.append(train_acc)\n        \n        if valid_loader is not None:\n            gc.collect()\n            para_valid_loader = pl.ParallelLoader(valid_loader, [device])\n            xm.master_print(f\"EPOCH {epoch} - VALIDATING...\")\n\n            with torch.no_grad():\n                valid_loss, valid_acc = model.validate_one_epoch(\n                    para_valid_loader.per_device_loader(device), criterion, device\n                )\n                \n            xm.master_print(f\"\\t[VALID] LOSS: {valid_loss:.4f}, ACCURACY: {valid_acc:.4f}\\n\")\n            valid_losses.append(valid_loss)\n            valid_accs.append(valid_acc)\n\n            if valid_loss < valid_loss_min:\n                xm.master_print(f\"Validation loss improved ({valid_loss_min:.4f} --> {valid_loss:.4f}). Saving model...\")\n                xm.save(model.state_dict(), f'best_model_epoch_{epoch}.pth')\n                valid_loss_min = valid_loss\n\n        xm.mark_step()\n        gc.collect()\n\n    return {\n        \"train_loss\": train_losses,\n        \"valid_loss\": valid_losses,\n        \"train_acc\": train_accs,\n        \"valid_acc\": valid_accs,\n    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:14.337656Z","iopub.execute_input":"2025-07-22T10:58:14.337883Z","iopub.status.idle":"2025-07-22T10:58:14.353266Z","shell.execute_reply.started":"2025-07-22T10:58:14.337862Z","shell.execute_reply":"2025-07-22T10:58:14.348009Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### `_run`\n\nEl código a continuación define la función `_run`, que se encarga de configurar y ejecutar el entrenamiento del modelo Vision Transformer (ViT) en un entorno distribuido utilizando TPUs (Tensor Processing Units). Esta función realiza varias tareas importantes, como la preparación de los datos, la inicialización del modelo, el entrenamiento y la validación, y finalmente guarda el modelo entrenado.\n\n1. **Cargar los Conjuntos de Datos**:\n   - Los datasets de entrenamiento (`train_df`) y validación (`valid_df`) se cargan utilizando la clase `CassavaDataset`, con las transformaciones especificadas en `transforms_train` y `transforms_valid`, respectivamente.\n\n2. **Samplers Distribuidos**:\n   - Se utilizan **samplers distribuidos** para asegurar que el entrenamiento y la validación se distribuyan correctamente entre las distintas TPUs. Esto es esencial cuando se entrena el modelo en varios dispositivos en paralelo.\n   - Los samplers distribuidos aseguran que cada TPU procese solo una parte del conjunto de datos.\n\n3. **DataLoader**:\n   - Se crean los **DataLoaders** para cargar los datos de entrenamiento y validación en lotes. Estos se configuran con un tamaño de lote (`BATCH_SIZE`), utilizando los samplers para distribuir los datos y especificando el número de workers para la carga paralela de los datos.\n\n4. **Inicialización del Modelo**:\n   - La función de pérdida `CrossEntropyLoss` se utiliza, que es comúnmente utilizada para tareas de clasificación.\n   - El dispositivo (TPU) se define utilizando `xm.xla_device()`, asegurando que el modelo se ejecute en TPUs en lugar de GPUs o CPUs.\n\n5. **Configuración del Optimizer y la Tasa de Aprendizaje**:\n   - Se configura un optimizador **Adam** con una tasa de aprendizaje ajustada según el número de TPUs. Esto es importante cuando se entrena en un entorno distribuido para garantizar que la tasa de aprendizaje sea adecuada.\n\n6. **Entrenamiento y Validación**:\n   - El entrenamiento se realiza utilizando la función `fit_tpu`, que se encarga de ejecutar el ciclo de entrenamiento y validación. Durante el entrenamiento, se calculan las métricas de pérdida y precisión y se imprimen para monitorear el progreso.\n   - Después de cada época de validación, se guarda el modelo si la pérdida de validación ha disminuido, asegurando que el modelo con el mejor rendimiento se conserve.\n\n7. **Guardado del Modelo**:\n   - Al final del entrenamiento, el modelo entrenado se guarda en un archivo con un nombre que incluye la fecha y hora actual para facilitar la organización de los modelos entrenados.\n","metadata":{}},{"cell_type":"code","source":"def _run():\n    import torch_xla.core.xla_model as xm\n    \n    try:\n        # Inicializar el dispositivo TPU con una mejor gestión de errores\n        device = xm.xla_device()\n        xm.master_print(f\"Successfully initialized TPU device: {device}\")\n        xm.master_print(f\"World size: {xm.xrt_world_size()}\")\n        xm.master_print(f\"Local rank: {xm.get_ordinal()}\")\n        \n    except Exception as e:\n        print(f\"Failed to initialize TPU: {e}\")\n        print(\"Falling back to CPU...\")\n        device = torch.device('cpu')\n    \n    # Crear el modelo\n    model = ViTBase16(n_classes=5, pretrained=False)\n    model.to(device)\n    \n    # crear los datasets\n    train_dataset = CassavaDataset(train_df, transforms=transforms_train)\n    valid_dataset = CassavaDataset(valid_df, transforms=transforms_valid)\n\n    # Comprobar si estamos utilizando TPU formación distribuida\n    if str(device).startswith('xla'):\n        train_sampler = DistributedSampler(\n            train_dataset,\n            num_replicas=xm.xrt_world_size(),\n            rank=xm.get_ordinal(),\n            shuffle=True,\n        )\n\n        valid_sampler = DistributedSampler(\n            valid_dataset,\n            num_replicas=xm.xrt_world_size(),\n            rank=xm.get_ordinal(),\n            shuffle=False,\n        )\n        \n        lr_multiplier = xm.xrt_world_size()\n    else:\n        train_sampler = None\n        valid_sampler = None\n        lr_multiplier = 1\n\n    # DataLoaders\n    train_loader = DataLoader(\n        dataset=train_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=train_sampler,\n        shuffle=(train_sampler is None),\n        drop_last=True,\n        num_workers=2,  \n    )\n\n    valid_loader = DataLoader(\n        dataset=valid_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=valid_sampler,\n        shuffle=False,\n        drop_last=True,\n        num_workers=2,\n    )\n\n    # Optimizador y criterios\n    criterion = nn.CrossEntropyLoss()\n    lr = LR * lr_multiplier\n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n\n    # Informacion de entrenamiento\n    print(f\"Device: {device}\")\n    print(f\"Learning Rate: {lr}\")\n    start_time = datetime.now()\n    print(f\"Start Time: {start_time}\")\n\n    # Entrenar\n    logs = fit_tpu(\n        model=model,\n        epochs=N_EPOCHS,\n        device=device,\n        criterion=criterion,\n        optimizer=optimizer,\n        train_loader=train_loader,\n        valid_loader=valid_loader,\n    )\n\n    # Guardar modelo final\n    end_time = datetime.now()\n    print(f\"Training completed in: {end_time - start_time}\")\n    \n    final_model_path = f'model_final_{end_time.strftime(\"%Y%m%d_%H%M\")}.pth'\n    if str(device).startswith('xla'):\n        xm.save(model.state_dict(), final_model_path)\n        xm.master_print(f\"Final model saved as: {final_model_path}\")\n    else:\n        torch.save(model.state_dict(), final_model_path)\n        print(f\"Final model saved as: {final_model_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:14.355398Z","iopub.execute_input":"2025-07-22T10:58:14.356218Z","iopub.status.idle":"2025-07-22T10:58:14.372776Z","shell.execute_reply.started":"2025-07-22T10:58:14.356191Z","shell.execute_reply":"2025-07-22T10:58:14.367428Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _mp_fn(rank, flags):\n    torch.set_default_tensor_type(\"torch.FloatTensor\")\n    _run()\n\n\nif __name__ == \"__main__\":\n    print(\"Attempting TPU training...\")\n    \n    # TPU multiprocessing\n    try:\n        import torch_xla.distributed.xla_multiprocessing as xmp\n        FLAGS = {}\n        print(\"Starting TPU multiprocessing training...\")\n        xmp.spawn(_mp_fn, args=(FLAGS,), start_method=\"fork\")\n    except Exception as e:\n        print(f\"TPU multiprocessing failed: {e}\")\n        \n        # Opcion 2: single TPU core\n        try:\n            print(\"Trying single TPU core...\")\n            _run()\n        except Exception as e2:\n            print(f\"Single TPU failed: {e2}\")\n            \n            # Opcion 3: CPU/GPU\n            print(\"Falling back to CPU/GPU training...\")\n            _run()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T10:58:14.374427Z","iopub.execute_input":"2025-07-22T10:58:14.374673Z"}},"outputs":[],"execution_count":null}]}