Tweaks

benchmarks/accelerate_opt/main.py

@@ -127,7 +127,6 @@ class CustomInitProcessGroupKwargs(InitProcessGroupKwargs):
             rank=int(os.environ["RANK"]),
             world_size=int(os.environ["WORLD_SIZE"]),
         )
-        print(init_process_group_kwargs.backend)
 
         # Accelerator SUCK, it is impossible to make it use hccl
         # We can bypass Accelerator logic by initializing the group ourselves

benchmarks/llama/main.py

@@ -59,10 +59,8 @@ def huggingface_main(args, model, config):
     import transformers
     from transformers import LlamaForCausalLM, LlamaTokenizerFast
     from transformers.models.llama.configuration_llama import LlamaConfig
-    from voir.wrapper import DataloaderWrapper, Wrapper
     from datasets import load_dataset
-    import optimum.habana
-
+
     # Dataset here
     println("Dataset")
     dataset = load_dataset("wikitext", "wikitext-103-v1")
@@ -173,8 +171,11 @@ def main():
 
 
 if __name__ == "__main__":
+    from voir.wrapper import StopProgram
+    import traceback
     try:
         main()
-    except Exception as err:
-        # Habana likes to eat exceptions
-        print(err)
+    except StopProgram:
+        print("Early stopped")
+    except Exception:
+        traceback.print_exc()

benchmarks/stargan/prepare.py

@@ -0,0 +1,14 @@
+
+
+
+def download_celebA():
+    from datasets import load_dataset
+    dataset = load_dataset(
+        "student/celebA", 
+        revision="2d31e6555722815c74ea7c845b07c1063dd705e9",
+        cache_dir="/tmp/milabench/cuda/results/data"
+    )
+
+
+if __name__ == "__main__":
+    download_celebA()

benchmarks/stargan/stargan/data_loader.py

@@ -91,6 +91,7 @@ def get_loader(
 
     if dataset == "CelebA":
         dataset = CelebA(image_dir, attr_path, selected_attrs, transform, mode)
+
     elif dataset == "RaFD":
         dataset = ImageFolder(image_dir, transform)
 

benchmarks/stargan/stargan/main.py

@@ -197,8 +197,13 @@ def main(config):
     parser.add_argument("--mode", type=str, default="train", choices=["train", "test"])
     parser.add_argument("--use_tensorboard", type=str2bool, default=False)
 
-    mbconfig = json.loads(os.environ["MILABENCH_CONFIG"])
-    datadir = mbconfig["dirs"]["extra"]
+    # try:
+    #     mbconfig = json.loads(os.environ["MILABENCH_CONFIG"])
+    #     datadir = mbconfig["dirs"]["extra"]
+    # except:
+    #     pass
+
+    datadir = "/tmp/milabench/cuda/results/data"
 
     # Directories.
     parser.add_argument("--celeba_image_dir", type=str, default="data/celeba/images")

benchmate/benchmate/dataset.py

@@ -0,0 +1,32 @@
+
+
+
+
+def no_transform(args):
+    return args
+
+
+def transform_images(transform_x, transform_y=no_transform):
+    def _(args):
+        print(args)
+        return transform_x(args[0]), transform_y(args[1])
+    return _
+
+
+def transform_celebA(transform_x):
+    def _(args):
+        print(args)
+        return transform_x(args["image"])
+    return _
+
+
+class TransformedDataset:
+    def __init__(self, dataset, transforms=no_transform):
+        self.dataset = dataset
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, item):
+        return self.transforms(self.dataset[item])

config/base.yaml

@@ -496,7 +496,8 @@ stargan:
     --image_size: 512
     --c_dim: 5
     --batch_size: 16
-    --dataset: "synth"
+    --dataset: "CelebA"
+    --celeba_image_dir: "{milabench_data}"
 
 super-slomo:
   inherits: _defaults

scripts/article/run_cuda.sh

@@ -53,6 +53,8 @@ install_prepare() {
     milabench prepare "$@"
 }
 
+module load cuda/12.3.2
+
 if [ ! -d "$MILABENCH_WORDIR" ]; then
     install_prepare 
 else