op and param counting

kaftanski · kaftanski · commit 8c9fca3854ec · 2023-01-05T13:55:20.000+01:00
diff --git a/.run/DGCNN seg train [master].run.xml b/.run/DGCNN seg train [master].run.xml
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="$PROJECT_DIR$/train.py" />
-    <option name="PARAMETERS" value="--data fissures --ds ts --pts 2048 --k 40 --static --coords --batch 32 --gpu 1 --output results/DGCNN_seg_cnn_image --kp_mode cnn --patch image --speed --offline" />
+    <option name="PARAMETERS" value="--data fissures --ds ts --pts 2048 --k 40 --static --coords --batch 32 --gpu 1 --output results/PointNet_seg_cnn_image --kp_mode cnn --patch image --test_only" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/.run/test DSEG-AE [master].run.xml b/.run/test DSEG-AE [master].run.xml
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="$PROJECT_DIR$/test_ae_regularization.py" />
-    <option name="PARAMETERS" value="--ds ts --gpu 3 --output results/DSEGAE_reg_aug_1024_foerstner_mind_ssc --seg_dir results/DGCNN_seg_foerstner_mind_ssc --ae_dir results/PC_AE_regularized_augment_1024 --offline" />
+    <option name="PARAMETERS" value="--ds ts --gpu 3 --output results/DSEGAE_reg_aug_1024_foerstner_mind_ssc --seg_dir results/DGCNN_seg_foerstner_mind_ssc --ae_dir results/PC_AE_regularized_augment_1024_folding --show" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/.run/train DG-SSM.run.xml b/.run/train DG-SSM.run.xml
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="$PROJECT_DIR$/train_dgcnn_ssm.py" />
-    <option name="PARAMETERS" value="--data fissures --ds ts --pts 2048 --k 40 --batch 8 --static --coords --gpu 1 --output results/DG-SSM_cnn_image --epochs 1000 --kp_mode cnn --patch image --predict_affine --scheduler cosine --head_schedule {\&quot;main\&quot;:90,\&quot;translation\&quot;:0,\&quot;rotation\&quot;:60,\&quot;scaling\&quot;:30} --offline" />
+    <option name="PARAMETERS" value="--data fissures --ds ts --pts 2048 --k 40 --batch 8 --static --coords --gpu 1 --output results/DG-SSM_cnn_image --epochs 1000 --kp_mode cnn --patch image --predict_affine --scheduler cosine --head_schedule {\&quot;main\&quot;:90,\&quot;translation\&quot;:0,\&quot;rotation\&quot;:60,\&quot;scaling\&quot;:30} --test_only" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/.run/train_pc_ae.run.xml b/.run/train_pc_ae.run.xml
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="$PROJECT_DIR$/train_pc_ae.py" />
-    <option name="PARAMETERS" value="--ds ts --k 20 --pts 1024 --batch 8 --output results/PC_AE_noreg_augment_1024 --loss_weights 1. 0. 0. 0. --shape plane --latent 512 --epochs 333 --deform --scheduler cosine --mesh --gpu 3 --offline" />
+    <option name="PARAMETERS" value="--ds ts --k 20 --pts 1024 --batch 8 --output results/PC_AE_regularized_augment_1024 --shape plane --latent 512 --epochs 333 --deform --scheduler cosine --mesh --gpu 3 --test_only" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/.run/train_segmentation_net.run.xml b/.run/train_segmentation_net.run.xml
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="$PROJECT_DIR$/train_segmentation_net.py" />
-    <option name="PARAMETERS" value="--data fissures --ds ts --batch 8 --gpu 1 --output results/lraspp_recall_loss --split ../nnUNet_baseline/nnu_preprocessed/Task503_FissuresTotalSeg/splits_final.pkl --epochs 50 --loss recall --patch_size 128 --model v3 --offline" />
+    <option name="PARAMETERS" value="--data fissures --ds ts --batch 8 --gpu 1 --output results/lraspp_recall_loss --split ../nnUNet_baseline/nnu_preprocessed/Task503_FissuresTotalSeg/splits_final.pkl --epochs 50 --loss recall --patch_size 128 --model v3 --test_only" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/thesis/utils.py b/thesis/utils.py
@@ -1,5 +1,9 @@
+import csv
 import os.path
 
+import torch
+from thop import profile, clever_format
+
 import matplotlib.pyplot as plt
 from matplotlib.patches import Patch
 
@@ -10,6 +14,19 @@
 plt.style.use('seaborn-paper')
 
 
+def param_and_op_count(model, input_shape, out_dir=None):
+    input = torch.zeros(input_shape)
+    macs, params = profile(model, (input, ))
+    if out_dir is not None:
+        with open(os.path.join(out_dir, 'op_count.csv'), 'w') as csv_file:
+            writer = csv.writer(csv_file)
+            writer.writerow(['Parameters', 'FLOPs'])
+            writer.writerow([params, macs])
+    macs, params = clever_format([macs, params], "%.3f")
+    print(macs, params)
+    return macs, params
+
+
 def save_fig(fig, outdir, basename_without_extension, dpi=300, show=True, pdf=True, padding=False):
     if not os.path.isdir(outdir):
         os.makedirs(outdir, exist_ok=True)
diff --git a/train.py b/train.py
@@ -20,6 +20,7 @@
 from metrics import assd, label_mesh_assd, batch_dice
 from models.dgcnn import DGCNNSeg
 from models.point_net import PointNetSeg
+from thesis.utils import param_and_op_count
 from utils.detached_run import maybe_run_detached_cli
 from utils.fissure_utils import binary_to_fissure_segmentation
 from utils.general_utils import kpts_to_world, mask_out_verts_from_mesh, remove_all_but_biggest_component, mask_to_points, \
@@ -616,6 +617,8 @@ def wrapped_fn(*args, **kwargs):
     net = model_class(in_features=in_features, num_classes=ds.num_classes, k=args.k,
                       spatial_transformer=args.transformer, dynamic=not args.static)
 
+    param_and_op_count(net, (1, *ds[0][0].shape), out_dir=args.output)
+
     if not args.test_only:
         store_args(args=args, out_dir=args.output)
 
diff --git a/train_dgcnn_ssm.py b/train_dgcnn_ssm.py
@@ -13,6 +13,7 @@
 from models.dg_ssm import DGSSM
 from shape_model.qualitative_evaluation import mode_plot
 from shape_model.ssm import vector2shape
+from thesis.utils import param_and_op_count
 from train import run, write_results, write_speed_results
 from utils.detached_run import maybe_run_detached_cli
 from utils.general_utils import no_print, get_device
@@ -202,6 +203,8 @@ def speed_test(ds, device, out_dir):
                   ssm_alpha=args.alpha, ssm_targ_var=args.target_variance, lssm=args.lssm,
                   predict_affine_params=args.predict_affine, only_affine=args.only_affine)
 
+    # param_and_op_count(DGSSM.load(args.output+'/fold0/model.pth', 'cpu'), (1, *ds[0][0].shape), out_dir=args.output)
+
     if not args.predict_affine:
         # set the loss weight for affine params to zero
         if args.loss_weights is None:
diff --git a/train_pc_ae.py b/train_pc_ae.py
@@ -14,6 +14,7 @@
 from data import CustomDataset
 from metrics import pseudo_symmetric_point_to_mesh_distance, assd
 from models.folding_net import DGCNNFoldingNet
+from thesis.utils import param_and_op_count
 from train import run, write_results
 from utils.detached_run import maybe_run_detached_cli
 from utils.image_ops import load_image_metadata
@@ -277,6 +278,8 @@ def test(ds: SampleFromMeshDS, device, out_dir, show):
     model = DGCNNFoldingNet(k=args.k, n_embedding=args.latent, shape_type=args.shape, decode_mesh=args.mesh,
                             deform=args.deform)
 
+    param_and_op_count(model, (ds.num_objects, *ds[0][0].shape), out_dir=args.output)
+
     # store config
     if not args.test_only:
         store_args(args=args, out_dir=args.output)
diff --git a/train_segmentation_net.py b/train_segmentation_net.py
@@ -13,6 +13,7 @@
 from metrics import batch_dice, binary_recall, binary_precision
 from models.lraspp_3d import LRASPP_MobileNetv3_large_3d
 from models.seg_cnn import MobileNetASPP
+from thesis.utils import param_and_op_count
 from train import run, write_results, compute_mesh_metrics
 from utils.detached_run import maybe_run_detached_cli
 from utils.fissure_utils import binary_to_fissure_segmentation
@@ -213,6 +214,8 @@ def test(ds: ImageDataset, device, out_dir, show):
     model_class = get_model_class(args)
     model = model_class(num_classes=ds.num_classes, patch_size=(args.patch_size,)*3)
 
+    param_and_op_count(model, (1, 1, *ds[0][0].shape), out_dir=args.output)
+
     # save setup
     if not args.test_only:
         store_args(args=args, out_dir=args.output)
