Subtle changes

.gitignore

@@ -8,4 +8,5 @@ work/StdContactTest/
 saved/MetalSet_*/
 saved/ViaSet_*/
 saved/StdMetal_*/
-saved/StdContact_*/
+saved/StdContact_*/
+dev/

lithobench/test.py

@@ -1,5 +1,6 @@
 import os
 import argparse
+from importlib import import_module
 
 from model import *
 from dataset import *
@@ -13,9 +14,13 @@
 from litho.damolitho import DAMOLitho
 from litho.cfnolitho import CFNOLitho
 
+# Example: python3 lithobench/test.py -m lithobench/ilt/neuralilt.py -a NeuralILT -i 512 -t ILT -o dev -s MetalSet -l saved/MetalSet_NeuralILT/net.pth
 def parseArgs(): 
     parser = argparse.ArgumentParser(description="Training ILT or Litho models")
-    parser.add_argument("--model", "-m", default="GANOPC", type=str, help="Model Name: {GANOPC, NeuralILT, DAMOILT, CFNOILT, LithoGAN, DOINN, DAMOLitho, CFNOLitho}")
+    parser.add_argument("--model", "-m", default="GANOPC", type=str, help="Model Name: {GANOPC, NeuralILT, DAMOILT, CFNOILT, LithoGAN, DOINN, DAMOLitho, CFNOLitho} or a user-provided filename")
+    parser.add_argument("--alias", "-a", default=None, type=str, help="Model alias, required when the model is provided by the user")
+    parser.add_argument("--task", "-t", default=None, type=str, help="Task: {ILT, Litho}")
+    parser.add_argument("--img_size", "-i", default=None, type=int, help="User-defined image size, required when the model is provided by the user")
     parser.add_argument("--benchmark", "-s", default="MetalSet", type=str, help="Benchmark: {MatelSet, ViaSet, StdMetal, StdContact}")
     parser.add_argument("--epochs", "-n", default=2, type=int, help="Number of epochs")
     parser.add_argument("--batch_size", "-b", default=4, type=int, help="Batch size")
@@ -37,13 +42,20 @@ def parseArgs():
         ImageSize = (256, 256)
     elif args.model in ["NeuralILT", ]: 
         ImageSize = (512, 512)
+    if not args.img_size is None: 
+        ImageSize = (args.img_size, args.img_size)
     Epochs = args.epochs
     BatchSize = args.batch_size
     NJobs = args.njobs
     Pretrain = args.pretrain
     EvalOnly = args.eval
     Task = "Litho" if args.model in ["LithoGAN", "DOINN", "DAMOLitho", "CFNOLitho"] else "ILT"
-    Folder = os.path.join(args.output, f"{args.benchmark}_{args.model}")
+    if not args.task is None: 
+        Task = args.task
+    alias = args.model
+    if not args.alias is None: 
+        alias = args.alias
+    Folder = os.path.join(args.output, f"{args.benchmark}_{alias}")
     if not os.path.exists(Folder): 
         os.mkdir(Folder)
     Filenames = None
@@ -75,7 +87,12 @@ def parseArgs():
     elif args.model == "CFNOLitho": 
         model = CFNOLitho(size=ImageSize)
     else: 
-        assert False, f"[ERROR]: Unsupported model: {args.model}"
+        if not os.path.exists(args.model): 
+            assert False, f"[ERROR]: Unsupported model: {args.model}"
+        modulename = ".".join(args.model[:-3].split("/"))
+        module = import_module(modulename)
+        MyModel = getattr(module, alias)
+        model = MyModel(size=ImageSize)
 
     if Benchmark in ["MetalSet", "ViaSet"]: 
         train_loader, val_loader = None, None

lithobench/train.py

@@ -1,5 +1,6 @@
 import os
 import argparse
+from importlib import import_module
 
 from model import *
 from dataset import *
@@ -13,9 +14,14 @@
 from litho.damolitho import DAMOLitho
 from litho.cfnolitho import CFNOLitho
 
+# Example: python3 lithobench/train.py -m lithobench/ilt/neuralilt.py -a NeuralILT -i 512 -t ILT -o dev -s MetalSet -n 8 -b 12 -p True
+#          python3 lithobench/train.py -m lithobench/litho/doinn.py -a DOINN -i 1024 -t Litho -o dev -s MetalSet -n 32 -b 16 -p False
 def parseArgs(): 
     parser = argparse.ArgumentParser(description="Training ILT or Litho models")
-    parser.add_argument("--model", "-m", default="GANOPC", type=str, help="Model Name: {GANOPC, NeuralILT, DAMOILT, CFNOILT, LithoGAN, DOINN, DAMOLitho, CFNOLitho}")
+    parser.add_argument("--model", "-m", default="GANOPC", type=str, help="Model name: {GANOPC, NeuralILT, DAMOILT, CFNOILT, LithoGAN, DOINN, DAMOLitho, CFNOLitho} or a user-provided filename")
+    parser.add_argument("--alias", "-a", default=None, type=str, help="Model alias, required when the model is provided by the user")
+    parser.add_argument("--task", "-t", default=None, type=str, help="Task: {ILT, Litho}")
+    parser.add_argument("--img_size", "-i", default=None, type=int, help="User-defined image size, required when the model is provided by the user")
     parser.add_argument("--benchmark", "-s", default="MetalSet", type=str, help="Benchmark: {MatelSet, ViaSet}")
     parser.add_argument("--epochs", "-n", default=2, type=int, help="Number of epochs")
     parser.add_argument("--batch_size", "-b", default=4, type=int, help="Batch size")
@@ -37,13 +43,20 @@ def parseArgs():
         ImageSize = (256, 256)
     elif args.model in ["NeuralILT", ]: 
         ImageSize = (512, 512)
+    if not args.img_size is None: 
+        ImageSize = (args.img_size, args.img_size)
     Epochs = args.epochs
     BatchSize = args.batch_size
     NJobs = args.njobs
     Pretrain = args.pretrain
     EvalOnly = args.eval
     Task = "Litho" if args.model in ["LithoGAN", "DOINN", "DAMOLitho", "CFNOLitho"] else "ILT"
-    Folder = os.path.join(args.output, f"{args.benchmark}_{args.model}")
+    if not args.task is None: 
+        Task = args.task
+    alias = args.model
+    if not args.alias is None: 
+        alias = args.alias
+    Folder = os.path.join(args.output, f"{args.benchmark}_{alias}")
     if not os.path.exists(Folder): 
         os.mkdir(Folder)
     Filenames = None
@@ -75,7 +88,12 @@ def parseArgs():
     elif args.model == "CFNOLitho": 
         model = CFNOLitho(size=ImageSize)
     else: 
-        assert False, f"[ERROR]: Unsupported model: {args.model}"
+        if not os.path.exists(args.model): 
+            assert False, f"[ERROR]: Unsupported model: {args.model}"
+        modulename = ".".join(args.model[:-3].split("/"))
+        module = import_module(modulename)
+        MyModel = getattr(module, alias)
+        model = MyModel(size=ImageSize)
 
     train_loader, val_loader = None, None
     targets = None

scripts/runCFNOLitho.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
-# python3 lithobench/train.py -m CFNOLitho -s MetalSet -n 8 -b 4 -p False
-python3 lithobench/train.py -m CFNOLitho -s ViaSet -n 2 -b 4 -p False
+python3 lithobench/train.py -m CFNOLitho -s MetalSet -n 8 -b 4 -p False
+# python3 lithobench/train.py -m CFNOLitho -s ViaSet -n 2 -b 4 -p False
 
 # python3 lithobench/test.py -m CFNOLitho -s MetalSet -b 4 -l saved/MetalSet_CFNOLitho/net.pth
 # python3 lithobench/test.py -m CFNOLitho -s ViaSet -b 4 -l saved/ViaSet_CFNOLitho/net.pth

scripts/runDAMOLitho.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
-# python3 lithobench/train.py -m DAMOLitho -s MetalSet -n 4 -b 4 -p True
-python3 lithobench/train.py -m DAMOLitho -s ViaSet -n 1 -b 4 -p True
+python3 lithobench/train.py -m DAMOLitho -s MetalSet -n 4 -b 4 -p True
+# python3 lithobench/train.py -m DAMOLitho -s ViaSet -n 1 -b 4 -p True
 
 # python3 lithobench/test.py -m DAMOLitho -s MetalSet -b 4 -g saved/MetalSet_DAMOLitho/netG.pth -d saved/MetalSet_DAMOLitho/netD.pth
 # python3 lithobench/test.py -m DAMOLitho -s ViaSe -b 4t -g saved/ViaSet_DAMOLitho/netG.pth -d saved/ViaSet_DAMOLitho/netD.pth

scripts/runDOINN.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
-# python3 lithobench/train.py -m DOINN -s MetalSet -n 32 -b 16 -p False
-python3 lithobench/train.py -m DOINN -s ViaSet -n 8 -b 16 -p False
+python3 lithobench/train.py -m DOINN -s MetalSet -n 32 -b 16 -p False
+# python3 lithobench/train.py -m DOINN -s ViaSet -n 8 -b 16 -p False
 
 # python3 lithobench/test.py -m DOINN -s MetalSet -b 16 -l saved/MetalSet_DOINN/net.pth
 # python3 lithobench/test.py -m DOINN -s ViaSet -b 16 -l saved/ViaSet_DOINN/net.pth