Fixed file descriptor leak

JonathanSeriesX · Oct 30, 2023 · c4ebe23 · c4ebe23
1 parent c6eccec
commit c4ebe23
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/data_loader.py b/data_loader.py
@@ -83,4 +83,7 @@ def __getitem__(self, idx):
         if self.transform:
             sample = self.transform(sample)
 
+        image.close()
+        label.close()
+
         return sample
diff --git a/u2net_train.py b/u2net_train.py
@@ -67,9 +67,11 @@ def get_args():
 def get_device():
     if torch.cuda.is_available():
         print("NVIDIA CUDA acceleration enabled")
+        torch.multiprocessing.set_start_method("spawn")
         return torch.device("cuda:0")
     elif torch.backends.mps.is_available():
         print("Apple Metal Performance Shaders acceleration enabled")
+        torch.multiprocessing.set_start_method("fork")
         return torch.device("mps")
     else:
         print("No GPU acceleration :/")
@@ -197,6 +199,8 @@ def main():
     args = get_args()
     device = get_device()
 
+    # subprocess.run(["ulimit", "-n", "4096"], shell=True, check=True)
+
     tra_image_dir = args.tra_image_dir
     tra_label_dir = args.tra_label_dir
     epoch_num = args.epoch_num