change inplace=True in ReLU

nidhi23ug · Mar 21, 2021 · 00f3675 · 00f3675
1 parent 48f6801
commit 00f3675
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 3 deletions.
diff --git a/train_classification.py b/train_classification.py
@@ -43,6 +43,12 @@ def parse_args():
     return parser.parse_args()
 
 
+def inplace_relu(m):
+    classname = m.__class__.__name__
+    if classname.find('ReLU') != -1:
+        m.inplace=True
+
+
 def test(model, loader, num_class=40):
     mean_correct = []
     class_acc = np.zeros((num_class, 3))
@@ -126,6 +132,7 @@ def log_string(str):
 
     classifier = model.get_model(num_class, normal_channel=args.use_normals)
     criterion = model.get_loss()
+    classifier.apply(inplace_relu)
 
     if not args.use_cpu:
         classifier = classifier.cuda()

diff --git a/train_partseg.py b/train_partseg.py
@@ -4,18 +4,19 @@
 """
 import argparse
 import os
-from data_utils.ShapeNetDataLoader import PartNormalDataset
 import torch
 import datetime
 import logging
-from pathlib import Path
 import sys
 import importlib
 import shutil
-from tqdm import tqdm
 import provider
 import numpy as np
 
+from pathlib import Path
+from tqdm import tqdm
+from data_utils.ShapeNetDataLoader import PartNormalDataset
+
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 ROOT_DIR = BASE_DIR
 sys.path.append(os.path.join(ROOT_DIR, 'models'))
@@ -30,6 +31,11 @@
         seg_label_to_cat[label] = cat
 
 
+def inplace_relu(m):
+    classname = m.__class__.__name__
+    if classname.find('ReLU') != -1:
+        m.inplace=True
+
 def to_categorical(y, num_classes):
     """ 1-hot encodes a tensor """
     new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
@@ -111,6 +117,7 @@ def log_string(str):
 
     classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda()
     criterion = MODEL.get_loss().cuda()
+    classifier.apply(inplace_relu)
 
     def weights_init(m):
         classname = m.__class__.__name__

diff --git a/train_semseg.py b/train_semseg.py
@@ -29,6 +29,10 @@
 for i, cat in enumerate(seg_classes.keys()):
     seg_label_to_cat[i] = cat
 
+def inplace_relu(m):
+    classname = m.__class__.__name__
+    if classname.find('ReLU') != -1:
+        m.inplace=True
 
 def parse_args():
     parser = argparse.ArgumentParser('Model')
@@ -111,6 +115,7 @@ def log_string(str):
 
     classifier = MODEL.get_model(NUM_CLASSES).cuda()
     criterion = MODEL.get_loss().cuda()
+    classifier.apply(inplace_relu)
 
     def weights_init(m):
         classname = m.__class__.__name__