Let's work with higher-res images then

JonathanSeriesX · Nov 1, 2023 · 409e62b · 409e62b
1 parent 74e60c0
commit 409e62b
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/u2net_train.py b/u2net_train.py
@@ -39,8 +39,8 @@
     "plain_resized": {
         "name": "Plain Images",
         "message": "Learning the dataset itself...\n",
-        "transform": [Resize(512), ToTensorLab()],
-        "batch_factor": 1,
+        "transform": [Resize(1024), ToTensorLab()],
+        "batch_factor": 0.3,
     },
     "flipped_v": {
         "name": "Vertical Flips",
@@ -367,13 +367,11 @@ def get_dataloader(tra_img_name_list, tra_lbl_name_list, transform, batch_size):
         transform=transform,
     )
 
-    cores = 8
-    if batch_size == 10:
-        cores = 2  # freeing up memory a bit
+    cores = 2  # freeing up memory a bit
 
     # DataLoader for the dataset
     dataloader = DataLoader(
-        dataset, batch_size=batch_size, shuffle=True, num_workers=cores
+        dataset, batch_size=int(batch_size), shuffle=True, num_workers=cores
     )
 
     return dataloader