fixed dataset load

Dpbm · Dpbm · commit a57111aef0bc · 2025-06-10T16:02:52.000-03:00
diff --git a/airflow-setup-test.sh b/airflow-setup-test.sh
@@ -1,14 +1,25 @@
 #!/bin/bash
 
+KAGGLE_KEY_FILE="$HOME/.kaggle/kaggle.json"
+
+clear_quotation_marks(){
+    STRING=$1
+    echo $STRING | sed 's/"//g'
+}
+
 export AIRFLOW_USERNAME="test"
 export AIRFLOW_PASSWORD="test"
 export AIRFLOW_EMAIL="test@test.com"
+
 export ROOT_DB="rootuser"
 export ROOT_PASS="rootpass"
 export DB_NAME="airflow"
 export DB_USERNAME="airflowuser"
 export DB_PASSWORD="airflowpass"
 
+export KAGGLE_USERNAME=$(clear_quotation_marks $(cat $KAGGLE_KEY_FILE | jq .username))
+export KAGGLE_KEY=$(clear_quotation_marks $(cat $KAGGLE_KEY_FILE | jq .key))
+
 # fix permission for volume
 sudo rm -rf ./data
 sudo mkdir -p ./data
diff --git a/airflow.Dockerfile b/airflow.Dockerfile
@@ -7,8 +7,6 @@ ENV PIPENV="/proj-venv/bin/pip"
 
 RUN ${PIPENV} install -r requirements.txt
 
-
-
 FROM apache/airflow:slim-2.11.0-python3.12 AS serve
 
 COPY --from=setup /proj-venv/lib/python3.12/site-packages/ /home/airflow/.local/lib/python3.12/site-packages
@@ -28,5 +26,7 @@ RUN airflow users create --username $USER \
     --email ${EMAIL} \
     --password ${PASSWORD}
 
+RUN airflow scheduler &
+
 EXPOSE 8080
-ENTRYPOINT [ "airflow", "standalone" ]
+ENTRYPOINT [ "airflow", "webserver", "--port", "8080" ]
diff --git a/airflow.dockerignore b/airflow.dockerignore
diff --git a/compose.yml b/compose.yml
@@ -40,4 +40,5 @@ services:
       - AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS=False
       - AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgres+psycopg2://${DB_USERNAME}:${DB_PASSWORD}@db:5432/${DB_NAME}
       - AIRFLOW_CONN_METADATA_DB=postgres+psycopg2://${DB_USERNAME}:${DB_PASSWORD}@db:5432/${DB_NAME}
-    restart: unless-stopped
+    restart: always
+
diff --git a/constants.py b/constants.py
@@ -30,3 +30,6 @@
 
 NEW_DIM = (500, 500)
 
+TRAIN_PERCENTAGE = 0.7
+TEST_PERCENTAGE = 0.2
+EVAL_PERCENTAGE = 1 - (TRAIN_PERCENTAGE + TEST_PERCENTAGE)
diff --git a/dags/dataset.py b/dags/dataset.py
@@ -3,6 +3,7 @@
 from airflow import DAG
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.operators.python import PythonOperator
+from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
 
 from dataset import generate_images, crate_dataset_folder, remove_duplicated_files, transform_images, start_df
 from constants import DATASET_PATH, DATASET_FILE
@@ -86,9 +87,21 @@
     Generate a GHZ experiment and saves the experiments results.
     """
 
+    trigger_dag_train = TriggerDagRunOperator(
+        task_id="run_training",
+        trigger_dag_id="train_model",
+        wait_for_completion=False
+    )
+
+    trigger_dag_train.doc_md = """
+    Run training after finishing all processes.
+    """
+
     create_folder >> [gen_ghz, gen_df]
     gen_df >> gen_images
     gen_images >> remove_duplicates
     remove_duplicates >> transform_img
     transform_img >> pack_img
 
+    [gen_ghz, pack_img] >> trigger_dag_train
+
diff --git a/dags/train.py b/dags/train.py
@@ -1,9 +1,6 @@
 """ETL pipeline using Airflow for training"""
 
-from datetime import timedelta
-
 from airflow import DAG
-from airflow.sensors.external_task import ExternalTaskSensor
 from airflow.providers.standard.operators.python import PythonOperator
 
 from train import setup_and_run_training
@@ -13,14 +10,6 @@
     description="train vision model"
 ) as dag:
 
-    wait_for_dataset_creation = ExternalTaskSensor(
-        task_id="wait_dataset_creation",
-        external_dag_id="build_dataset",
-        external_task_id=None,  # Set to None to wait for the whole DAG
-        mode="poke",
-    )
-
-
     train = PythonOperator(
         task_id="train_model",
         python_callable=setup_and_run_training,
@@ -29,5 +18,3 @@
     train.doc_md = """
     Run the training cycle
     """
-
-    wait_for_dataset_creation >> train
diff --git a/export/__init__.py b/export/__init__.py
diff --git a/export/huggingface.py b/export/huggingface.py
diff --git a/export/kaggle.py b/export/kaggle.py
@@ -0,0 +1,2 @@
+"""Export dataset and model to kaggle"""
+import kagglehub as kh
diff --git a/helpers.py b/helpers.py
@@ -32,8 +32,10 @@ def plot_filters(images:torch.Tensor, title:Optional[str]=None):
 
         plt.axis('off')
 
+
         for i,image in enumerate(images):
             ax = plt.subplot(rows,cols,i+1)
+            ax.axis('off')
             ax.imshow(image)
         plt.show()
 
diff --git a/train.py b/train.py
@@ -26,10 +26,14 @@
     DEBUG, 
     BATCH_SIZE,
     GHZ_FILE,
-    GHZ_PRED_FILE
+    GHZ_PRED_FILE,
+    TRAIN_PERCENTAGE,
+    TEST_PERCENTAGE,
+    EVAL_PERCENTAGE
 )
 from helpers import debug, PlotImages
 from colors import Colors
+from dataset import get_schema
 
 StateDict = OrderedDict
 Device = str
@@ -39,15 +43,24 @@
 class ImagesDataset(Dataset):
     """Dataset class for handling batches and data itself"""
 
-    def __init__(self, device:Device, file:FilePath, dataset_file:FilePath):
-        self._dataset = pl.read_csv(dataset_file)
-        self._obj = h5py.File(file, "r")
-        self._total = len(self._obj)
+    def __init__(
+        self, 
+        device:Device, 
+        file:h5py.File, 
+        dataset:pl.LazyFrame,
+        total_images:int,
+        pivot:int
+        ):
+
+        self._dataset = dataset
+        self._obj = file
+        self._pivot = pivot
         self._device = device
+        self._total_images = total_images
 
     def __len__(self) -> int:
         """return the amount of files"""
-        return self._total
+        return self._total_images
 
     def _to_tensor(self, loaded_file:np.array) -> torch.Tensor:
         """auxiliary method to map an np.array to tensor in the correct device and data type"""
@@ -58,8 +71,13 @@ def _to_tensor(self, loaded_file:np.array) -> torch.Tensor:
 
     def __getitem__(self, index:int) -> Tuple[torch.Tensor, torch.Tensor]:
         """Get an especific value inside the dataset with its label"""
-        input_data = self._to_tensor(self._obj[f"{index}"][()])
-        label = torch.from_numpy(np.array(json.loads(self._dataset.row(index, named=True)["result"]), dtype=np.float16)).to(self._device)
+        shifted_index = index+self._pivot
+        input_data = self._to_tensor(self._obj[f"{shifted_index}"][()])
+
+        dataset_data = self._dataset.slice(length=1, offset=index).collect()
+        result = json.loads(dataset_data["result"].item(0))
+
+        label = torch.from_numpy(np.array(result, dtype=np.float16)).to(self._device)
         return input_data, label
 
 class Downsample(torch.nn.Module):
@@ -118,7 +136,7 @@ def __init__(self):
         self.pool1 = nn.MaxPool2d(2, stride=2)
         self.pool2 = nn.AvgPool2d(3, stride=2)
 
-        self.out_neurons = 512*9*18
+        self.out_neurons = 512*6*7
         self.fc1 = nn.Linear(self.out_neurons, 32)
 
         self.blocks = nn.ModuleList([
@@ -282,11 +300,33 @@ def train(device:Device, checkpoint:Checkpoint):
 
     model = get_model(device, checkpoint.model)
 
-    train_data = ImagesDataset(device, IMAGES_TRAIN, DATASET_FILE)
-    test_data = ImagesDataset(device, IMAGES_TEST, DATASET_FILE)
+    dataset = pl.scan_csv(DATASET_FILE, schema=get_schema())
+
+    h5_file = h5py.File(IMAGES_H5_FILE, "r")
+    total_images = len(h5_file)
+
+    # the amount of images for TRAINING, TESTING AND EVALUATING
+    max_train = int(TRAIN_PERCENTAGE*total_images)
+    max_test = int(TEST_PERCENTAGE*total_images)
+    max_eval = int(EVAL_PERCENTAGE*total_images)
+
+    # the starting index for images per dataset split
+    pivot_training_index = 0
+    pivot_testing_index = max_train - 1
+    pivot_evaluating_index = total_images - max_eval - 1
+
+    # split dataset (csv file) into sections
+    train_dataset = dataset.slice(length=max_train, offset=pivot_training_index)
+    test_dataset = dataset.slice(length=max_test, offset=pivot_testing_index)
+    eval_dataset = dataset.slice(length=max_eval, offset=pivot_evaluating_index)
+
+    train_data = ImagesDataset(device, h5_file, train_dataset, max_train, pivot_training_index)
+    test_data = ImagesDataset(device, h5_file, test_dataset, max_test, pivot_testing_index)
+    eval_data = ImagesDataset(device, h5_file, eval_dataset, max_eval, pivot_evaluating_index)
     
     data_loader_train = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
     data_loader_test = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True)
+    data_loader_eval = DataLoader(eval_data, batch_size=BATCH_SIZE, shuffle=True)
 
     loss_fn = nn.KLDivLoss(reduction="batchmean")
     lr = 0.1
@@ -318,8 +358,8 @@ def train(device:Device, checkpoint:Checkpoint):
                loss = loss_fn(output.log(),label)
                running_loss += loss
 
-        avg_loss = running_loss/(i+1)
-        print("%sAVG: %f%s"%(Colors.MAGENTAFG, avg_loss, Colors.ENDC))
+        avg_loss = running_loss/max_test
+        print("%sAVG Test: %f%s"%(Colors.MAGENTAFG, avg_loss, Colors.ENDC))
 
         if avg_loss < best_loss:
             best_loss = avg_loss
@@ -334,6 +374,8 @@ def train(device:Device, checkpoint:Checkpoint):
             scheduler.state_dict()
         )
 
+    h5_file.close()
+
     return model
 
 def get_device() -> Device:
@@ -351,6 +393,10 @@ def get_device() -> Device:
 
     return device
 
+def get_RMSE(c):
+    # squared_diff_sum = torch.sqrt()
+
+
 def get_checkpoint_arg() -> Optional[str]:
     """
     Parse args from command line and retrieve the
@@ -372,9 +418,12 @@ def run_debug_experiemnt():
     checkpoint_path = get_checkpoint_arg()     
     checkpoint = Checkpoint(checkpoint_path)
     checkpoint.load()
+    
+    dataset = pl.scan_csv(DATASET_FILE, schema=get_schema())
+    h5_file = h5py.File(IMAGES_H5_FILE, "r")
 
     model = get_model(device, checkpoint.model)
-    test = ImagesDataset(device, IMAGES_TRAIN, DATASET_FILE)
+    test = ImagesDataset(device, h5_file, dataset, len(h5_file), 0)
     test_loader = DataLoader(test, batch_size=1, shuffle=True)
     model.train(False)
     model.eval()
@@ -398,6 +447,7 @@ def setup_and_run_training():
     checkpoint = Checkpoint(checkpoint_path)
     checkpoint.load()
 
+    # TODO: USE PYTORCH SCRIPTING (JIT)
     model = train(device,checkpoint)
     model.save() # save best model
     model.eval()
@@ -422,4 +472,4 @@ def main():
     try:
         main()
     except KeyboardInterrupt:
-        sys.exit(0)
+        sys.exit(0)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+"""Export dataset and model to kaggle"""`
	`2`	`+import kagglehub as kh`