pb de gpu

Author: rayaberova7

training/argo-workflows/train-workflow.yaml

@@ -16,14 +16,14 @@ spec:
         value: '[
             { "SOURCE": "SENTINEL2",
               "N_BANDS": "14",
-              "DATASETS": ["AT332_2018","BE100_2018","BE251_2018","BG322_2018","CY000_2018","CZ072_2018","DEA54_2018","DK041_2018","EE00A_2018","EL521_2018","ES612_2018","FI1C1_2018","FRJ27_2018","FRK26_2018","HR050_2018","IE061_2018","ITI32_2018","LT028_2018","LU000_2018","LV008_2018","MT001_2018","NL33C_2018","PL414_2018","PT16I_2018","RO123_2018","SI035_2018","SK022_2018","UKJ22_2018","AT332_2021","BE100_2021","BE251_2021","BG322_2021","CY000_2021","CZ072_2021","DEA54_2021","DK041_2021","EE00A_2021","EL521_2021","ES612_2021","FI1C1_2021","FRJ27_2021","FRK26_2021","HR050_2021","IE061_2021","ITI32_2021","LT028_2021","LU000_2021","LV008_2021","MT001_2021","NL33C_2021","PL414_2021","PT16I_2021","RO123_2021","SI035_2021","SK022_2021"],
+              "DATASETS": ["AT332_2018","BE100_2018","BE251_2018","BG322_2018","CY000_2018","CZ072_2018","DEA54_2018","DK041_2018","EE00A_2018","EL521_2018","ES612_2018","FI1C1_2018","FRJ27_2018","FRK26_2018","HR050_2018","IE061_2018","ITI32_2018","LT028_2018","LU000_2018","LV008_2018","MT001_2018","NL33C_2018","PL414_2018","PT16I_2018","RO123_2018","SI035_2018","SK022_2018","UKJ22_2018","AT332_2021","BE100_2021","BE251_2021","BG322_2021","CY000_2021","CZ072_2021","DEA54_2021","EE00A_2021","EL521_2021","ES612_2021","FI1C1_2021","FRJ27_2021","FRK26_2021","HR050_2021","IE061_2021","ITI32_2021","LT028_2021","LU000_2021","LV008_2021","MT001_2021","NL33C_2021","PL414_2021","PT16I_2021","RO123_2021","SI035_2021","SK022_2021"],
               "TILES_SIZE": 250,
               "AUGMENT_SIZE": 512,
               "TYPE_LABELER": "CLCplus-Backbone",
               "USE_S3": 0,
-              "EPOCHS": 1,
-              "BATCH_SIZE": 4,
-              "TEST_BATCH_SIZE": 2,
+              "EPOCHS": 50,
+              "BATCH_SIZE": 8,
+              "TEST_BATCH_SIZE": 4,
               "LR": 0.00005,
               "BUILDING_CLASS_WEIGHT": 10.0,
               "LOSS_NAME": "cross_entropy_weighted",
@@ -113,6 +113,8 @@ spec:
           - name: LOGITS
           - name: FREEZE_ENCODER
           - name: CUDA
+      nodeSelector:
+             gpu-vram: 80GB
       container:
         image: inseefrlab/onyxia-python-pytorch:py3.13.12-gpu
         imagePullPolicy: Always
@@ -183,4 +185,4 @@ spec:
                 name: mlflow
                 key: password
           - name: ENTRY_POINT
-            value: main
+            value: main

training/src/train.py

@@ -281,7 +281,7 @@ def main(
     random.seed(0)
     np.random.seed(0)
 
-    kwargs = {"num_workers": os.cpu_count(), "pin_memory": True} if args.cuda else {}
+    kwargs = {"num_workers": 0, "pin_memory": False} if args.cuda else {}
 
     earlystop = {"monitor": "val_loss", "patience": patience, "mode": "min"}
     checkpoints = [