Pointnet_Example/model.py

@@ -73,13 +73,13 @@ class PointNet(nn.Module):
         def forward(self, input):
             matrix3x3 = self.input_transform(input)
             # batch matrix multiplication
-            xb = torch.bmm(torch.transpose(input, 1, 2),
+            xb = torch.matmul(torch.transpose(input, 1, 2),
                            matrix3x3).transpose(1, 2)
 
             xb = F.relu(self.bn1(self.conv1(xb)))
 
             matrix64x64 = self.feature_transform(xb)
-            xb = torch.bmm(torch.transpose(xb, 1, 2),
+            xb = torch.matmul(torch.transpose(xb, 1, 2),
                            matrix64x64).transpose(1, 2)
 
             xb = F.relu(self.bn2(self.conv2(xb)))
@@ -104,8 +104,8 @@ def pointnetloss(outputs, labels, m3x3, m64x64, alpha=0.001, device=None):
     # Calculate difference to identity matrix for regularization.
     id3x3 = torch.eye(3, requires_grad=True, device=device).repeat(bs, 1, 1)
     id64x64 = torch.eye(64, requires_grad=True, device=device).repeat(bs, 1, 1)
-    diff3x3 = id3x3 - torch.bmm(m3x3, m3x3.transpose(1, 2))
-    diff64x64 = id64x64 - torch.bmm(m64x64, m64x64.transpose(1, 2))
+    diff3x3 = id3x3 - torch.matmul(m3x3, m3x3.transpose(1, 2))
+    diff64x64 = id64x64 - torch.matmul(m64x64, m64x64.transpose(1, 2))
     # Negative log likelihood criterion is already adapted to batch size.
     return criterion(outputs, labels) + alpha * (torch.norm(diff3x3) + torch.norm(diff64x64)) / float(bs)
 

Pointnet_Example/submit_train.sh

 #!/bin/bash
 #SBATCH --job-name=train-forest-script
 #SBATCH -p gpu                       # request gpu node for the training
-#SBATCH -t 00:05:00                  # TODO: estimate the time you will need
-#SBATCH -G rtx5000                   # requesting specific GPU, run sinfo -p gpu --format=%N,%G # to see what is available
+#SBATCH -t 00:15:00                  # TODO: estimate the time you will need
+#SBATCH -G gtx1080                   # requesting specific GPU, run sinfo -p gpu --format=%N,%G # to see what is available
 #SBATCH --nodes=1                    # total number of nodes
 #SBATCH --ntasks=1                   # total number of tasks
 #SBATCH --mail-type=begin            # send mail when job begins
@@ -21,7 +21,10 @@ echo "Home directory: ${HOME}"
 echo "Working directory: $PWD"
 echo "Current node: ${SLURM_NODELIST}"
 
+# For debugging purposes.
 python --version
+python -m torch.utils.collect_env
+nvcc -V
 
 # Run the script.
-python train.py
\ No newline at end of file
+python -u train.py
\ No newline at end of file

Pointnet_Example/utils.py

@@ -174,3 +174,11 @@ def remove_empty_las_files(start_dir: str, verbose: bool = True) -> None:
                 os.remove(full_file_name)
                 if verbose:
                     print(f"Remove {full_file_name}.")
+
+def print_memory():
+    """Print the GPU memory. Use for debugging."""
+    total = torch.cuda.get_device_properties(0).total_memory
+    reserved = torch.cuda.memory_reserved(0)
+    allocated = torch.cuda.memory_allocated(0)
+    free = reserved - allocated
+    print(f"Memory: Total {total} | Reserved {reserved} | Allocated {allocated} | Free memory {free}")
\ No newline at end of file