Adds initial 3D CNN voxel implementation

Implements a fast 3D CNN for vertex prediction from voxelized point cloud patches. It uses 3D convolutions and voxelization and includes data loading and training scripts. Also adds a batch script for cluster training.

fast_voxel.py

@@ -0,0 +1,591 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import pickle
+from torch.utils.data import Dataset, DataLoader
+from typing import List, Dict, Tuple, Optional
+import json
+
+class Fast3DCNN(nn.Module):
+    """
+    Fast 3D CNN implementation for 3D vertex prediction from voxelized point cloud patches.
+    Takes 7D point clouds (x,y,z,r,g,b,filtered_flag) and predicts 3D vertex coordinates.
+    Uses voxelization and 3D convolutions instead of PointNet architecture.
+    """
+    def __init__(self, input_channels=7, output_dim=3, voxel_size=32, predict_score=True, predict_class=True, num_classes=1):
+        super(Fast3DCNN, self).__init__()
+        self.voxel_size = voxel_size
+        self.predict_score = predict_score
+        self.predict_class = predict_class
+        self.num_classes = num_classes
+        
+        # 3D Convolutional layers for feature extraction
+        self.conv1 = nn.Conv3d(input_channels, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv3d(64, 128, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv3d(128, 256, kernel_size=3, padding=1)
+        self.conv4 = nn.Conv3d(256, 512, kernel_size=3, padding=1)
+        self.conv5 = nn.Conv3d(512, 512, kernel_size=3, padding=1)
+        
+        # Additional convolutional layers for deeper feature extraction
+        self.conv6 = nn.Conv3d(512, 1024, kernel_size=3, padding=1)
+        
+        # Batch normalization layers
+        self.bn1 = nn.BatchNorm3d(64)
+        self.bn2 = nn.BatchNorm3d(128)
+        self.bn3 = nn.BatchNorm3d(256)
+        self.bn4 = nn.BatchNorm3d(512)
+        self.bn5 = nn.BatchNorm3d(512)
+        self.bn6 = nn.BatchNorm3d(1024)
+        
+        # Max pooling layers
+        self.pool = nn.MaxPool3d(kernel_size=2, stride=2)
+        
+        # Calculate the size after convolutions and pooling
+        # Starting with voxel_size^3, after 3 pooling operations: voxel_size / 8
+        final_size = voxel_size // 8
+        flattened_size = 1024 * (final_size ** 3)
+        
+        # Adaptive pooling to handle variable sizes
+        self.adaptive_pool = nn.AdaptiveAvgPool3d((4, 4, 4))
+        flattened_size = 1024 * 4 * 4 * 4
+        
+        # Shared fully connected layers
+        self.shared_fc1 = nn.Linear(flattened_size, 1024)
+        self.shared_fc2 = nn.Linear(1024, 512)
+        
+        # Position prediction head
+        self.pos_fc1 = nn.Linear(512, 512)
+        self.pos_fc2 = nn.Linear(512, 256)
+        self.pos_fc3 = nn.Linear(256, 128)
+        self.pos_fc4 = nn.Linear(128, output_dim)
+        
+        # Score prediction head
+        if self.predict_score:
+            self.score_fc1 = nn.Linear(512, 512)
+            self.score_fc2 = nn.Linear(512, 256)
+            self.score_fc3 = nn.Linear(256, 128)
+            self.score_fc4 = nn.Linear(128, 64)
+            self.score_fc5 = nn.Linear(64, 1)
+        
+        # Classification head
+        if self.predict_class:
+            self.class_fc1 = nn.Linear(512, 512)
+            self.class_fc2 = nn.Linear(512, 256)
+            self.class_fc3 = nn.Linear(256, 128)
+            self.class_fc4 = nn.Linear(128, 64)
+            self.class_fc5 = nn.Linear(64, num_classes)
+        
+        # Dropout layers
+        self.dropout_light = nn.Dropout(0.2)
+        self.dropout_medium = nn.Dropout(0.3)
+        self.dropout_heavy = nn.Dropout(0.4)
+
+    def forward(self, x):
+        """
+        Forward pass
+        Args:
+            x: (batch_size, input_channels, voxel_size, voxel_size, voxel_size) tensor
+        Returns:
+            Tuple containing predictions based on configuration:
+            - position: (batch_size, output_dim) tensor of predicted 3D coordinates
+            - score: (batch_size, 1) tensor of predicted distance to GT (if predict_score=True)
+            - classification: (batch_size, num_classes) tensor of class logits (if predict_class=True)
+        """
+        batch_size = x.size(0)
+        
+        # 3D Convolutional feature extraction
+        x1 = F.relu(self.bn1(self.conv1(x)))
+        x1 = self.pool(x1)
+        
+        x2 = F.relu(self.bn2(self.conv2(x1)))
+        x2 = self.pool(x2)
+        
+        x3 = F.relu(self.bn3(self.conv3(x2)))
+        x3 = self.pool(x3)
+        
+        x4 = F.relu(self.bn4(self.conv4(x3)))
+        x5 = F.relu(self.bn5(self.conv5(x4)))
+        x6 = F.relu(self.bn6(self.conv6(x5)))
+        
+        # Adaptive pooling to ensure consistent size
+        x6 = self.adaptive_pool(x6)
+        
+        # Flatten for fully connected layers
+        global_features = x6.view(batch_size, -1)
+        
+        # Shared features
+        shared1 = F.relu(self.shared_fc1(global_features))
+        shared1 = self.dropout_light(shared1)
+        shared2 = F.relu(self.shared_fc2(shared1))
+        shared_features = self.dropout_medium(shared2)
+        
+        # Position prediction
+        pos1 = F.relu(self.pos_fc1(shared_features))
+        pos1 = self.dropout_light(pos1)
+        pos2 = F.relu(self.pos_fc2(pos1))
+        pos2 = self.dropout_medium(pos2)
+        pos3 = F.relu(self.pos_fc3(pos2))
+        pos3 = self.dropout_light(pos3)
+        position = self.pos_fc4(pos3)
+        
+        outputs = [position]
+        
+        if self.predict_score:
+            # Score prediction
+            score1 = F.relu(self.score_fc1(shared_features))
+            score1 = self.dropout_light(score1)
+            score2 = F.relu(self.score_fc2(score1))
+            score2 = self.dropout_medium(score2)
+            score3 = F.relu(self.score_fc3(score2))
+            score3 = self.dropout_light(score3)
+            score4 = F.relu(self.score_fc4(score3))
+            score4 = self.dropout_light(score4)
+            score = F.relu(self.score_fc5(score4))
+            outputs.append(score)
+        
+        if self.predict_class:
+            # Classification prediction
+            class1 = F.relu(self.class_fc1(shared_features))
+            class1 = self.dropout_light(class1)
+            class2 = F.relu(self.class_fc2(class1))
+            class2 = self.dropout_medium(class2)
+            class3 = F.relu(self.class_fc3(class2))
+            class3 = self.dropout_light(class3)
+            class4 = F.relu(self.class_fc4(class3))
+            class4 = self.dropout_light(class4)
+            classification = self.class_fc5(class4)
+            outputs.append(classification)
+        
+        # Return outputs based on configuration
+        if len(outputs) == 1:
+            return outputs[0]
+        elif len(outputs) == 2:
+            if self.predict_score:
+                return outputs[0], outputs[1]
+            else:
+                return outputs[0], outputs[1]
+        else:
+            return outputs[0], outputs[1], outputs[2]
+
+def voxelize_patch(patch_7d: np.ndarray, voxel_size: int = 32, patch_size: float = 1.0) -> np.ndarray:
+    """
+    Convert point cloud patch to voxel grid.
+    
+    Args:
+        patch_7d: (N, 7) array of points with [x, y, z, r, g, b, filtered_flag]
+        voxel_size: Size of the voxel grid (voxel_size^3)
+        patch_size: Physical size of the patch in world coordinates
+        
+    Returns:
+        voxels: (7, voxel_size, voxel_size, voxel_size) array of voxelized features
+    """
+    if len(patch_7d) == 0:
+        return np.zeros((7, voxel_size, voxel_size, voxel_size))
+    
+    # Extract coordinates and features
+    coords = patch_7d[:, :3]  # x, y, z
+    features = patch_7d[:, 3:]  # r, g, b, filtered_flag
+    
+    # Normalize coordinates to [0, voxel_size-1]
+    coords_min = coords.min(axis=0)
+    coords_max = coords.max(axis=0)
+    coords_range = coords_max - coords_min
+    coords_range[coords_range == 0] = 1  # Avoid division by zero
+    
+    normalized_coords = (coords - coords_min) / coords_range * (voxel_size - 1)
+    voxel_indices = normalized_coords.astype(int)
+    
+    # Clip to valid range
+    voxel_indices = np.clip(voxel_indices, 0, voxel_size - 1)
+    
+    # Initialize voxel grid
+    voxels = np.zeros((7, voxel_size, voxel_size, voxel_size))
+    
+    # Fill voxels with features (average if multiple points fall in same voxel)
+    counts = np.zeros((voxel_size, voxel_size, voxel_size))
+    
+    for i in range(len(patch_7d)):
+        x, y, z = voxel_indices[i]
+        # Store normalized coordinates in first 3 channels
+        voxels[0, x, y, z] += normalized_coords[i, 0] / (voxel_size - 1)  # normalized x
+        voxels[1, x, y, z] += normalized_coords[i, 1] / (voxel_size - 1)  # normalized y
+        voxels[2, x, y, z] += normalized_coords[i, 2] / (voxel_size - 1)  # normalized z
+        # Store RGB and filtered_flag in remaining channels
+        voxels[3:, x, y, z] += features[i]
+        counts[x, y, z] += 1
+    
+    # Average features where multiple points exist
+    mask = counts > 0
+    for c in range(7):
+        voxels[c][mask] /= counts[mask]
+    
+    return voxels
+
+class VoxelPatchDataset(Dataset):
+    """
+    Dataset class for loading saved patches and converting them to voxel grids for 3D CNN training.
+    """
+    
+    def __init__(self, dataset_dir: str, voxel_size: int = 32, augment: bool = False):
+        self.dataset_dir = dataset_dir
+        self.voxel_size = voxel_size
+        self.augment = augment
+        
+        # Load patch files
+        self.patch_files = []
+        for file in os.listdir(dataset_dir):
+            if file.endswith('.pkl'):
+                self.patch_files.append(os.path.join(dataset_dir, file))
+                
+        print(f"Found {len(self.patch_files)} patch files in {dataset_dir}")
+
+    def __len__(self):
+        return len(self.patch_files)
+
+    def __getitem__(self, idx):
+        """
+        Load and process a patch for training.
+        Returns:
+            voxel_data: (7, voxel_size, voxel_size, voxel_size) tensor of voxelized data
+            target: (3,) tensor of target 3D coordinates
+            valid_mask: scalar tensor indicating if this is a valid sample
+            distance_to_gt: scalar tensor of distance from initial prediction to GT
+            classification: scalar tensor for binary classification (1 if GT vertex present, 0 if not)
+        """
+        patch_file = self.patch_files[idx]
+        
+        with open(patch_file, 'rb') as f:
+            patch_info = pickle.load(f)
+        
+        patch_7d = patch_info['patch_7d']  # (N, 7)
+        target = patch_info.get('assigned_wf_vertex', None)  # (3,) or None
+        initial_pred = patch_info.get('cluster_center', None)  # (3,) or None
+        
+        # Determine classification label based on GT vertex presence
+        has_gt_vertex = 1.0 if target is not None else 0.0
+        
+        # Handle patches without ground truth
+        if target is None:
+            target = np.zeros(3)
+        else:
+            target = np.array(target)
+        
+        # Voxelize the patch
+        voxel_data = voxelize_patch(patch_7d, self.voxel_size)
+        
+        # Data augmentation (only if GT vertex is present)
+        if self.augment and has_gt_vertex > 0:
+            voxel_data, target = self._augment_voxels(voxel_data, target)
+        
+        # Convert to tensors (copy arrays to handle negative strides from augmentation)
+        voxel_tensor = torch.from_numpy(voxel_data.copy()).float()  # (7, voxel_size, voxel_size, voxel_size)
+        target_tensor = torch.from_numpy(target.copy()).float()  # (3,)
+        
+        # Valid mask (check if voxel grid has any non-zero values)
+        valid_mask = torch.tensor(1.0 if voxel_data.sum() > 0 else 0.0)
+        
+        # Handle initial_pred
+        if initial_pred is not None:
+            initial_pred_tensor = torch.from_numpy(initial_pred).float()
+        else:
+            initial_pred_tensor = torch.zeros(3).float()
+        
+        # Classification tensor
+        classification_tensor = torch.tensor(has_gt_vertex).float()
+        
+        return voxel_tensor, target_tensor, valid_mask, initial_pred_tensor, classification_tensor
+
+    def _augment_voxels(self, voxel_data: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Apply data augmentation to voxel data.
+        """
+        # Random rotation around Z-axis
+        if np.random.random() > 0.5:
+            k = np.random.randint(1, 4)  # 90, 180, or 270 degrees
+            voxel_data = np.rot90(voxel_data, k, axes=(1, 2))  # Rotate around z-axis
+        
+        # Random flip
+        if np.random.random() > 0.5:
+            voxel_data = np.flip(voxel_data, axis=1)  # Flip along x-axis
+        if np.random.random() > 0.5:
+            voxel_data = np.flip(voxel_data, axis=2)  # Flip along y-axis
+        
+        return voxel_data, target
+
+def save_patches_dataset(patches: List[Dict], dataset_dir: str, entry_id: str):
+    """
+    Save patches from prediction pipeline to create a training dataset.
+    
+    Args:
+        patches: List of patch dictionaries from generate_patches()
+        dataset_dir: Directory to save the dataset
+        entry_id: Unique identifier for this entry/image
+    """
+    os.makedirs(dataset_dir, exist_ok=True)
+    
+    for i, patch in enumerate(patches):
+        # Create unique filename
+        filename = f"{entry_id}_patch_{i}.pkl"
+        filepath = os.path.join(dataset_dir, filename)
+        
+        # Skip if file already exists
+        if os.path.exists(filepath):
+            continue
+        
+        # Save patch data
+        with open(filepath, 'wb') as f:
+            pickle.dump(patch, f)
+    
+    print(f"Saved {len(patches)} patches for entry {entry_id}")
+
+# Create dataloader with custom collate function to filter invalid samples
+def collate_fn(batch):
+    valid_batch = []
+    for voxel_data, target, valid_mask, initial_pred, classification in batch:
+        # Filter out invalid samples
+        if valid_mask > 0:
+            valid_batch.append((voxel_data, target, valid_mask, initial_pred, classification))
+    
+    if len(valid_batch) == 0:
+        return None
+    
+    # Stack valid samples
+    voxel_data = torch.stack([item[0] for item in valid_batch])
+    targets = torch.stack([item[1] for item in valid_batch])
+    valid_masks = torch.stack([item[2] for item in valid_batch])
+    initial_preds = torch.stack([item[3] for item in valid_batch])
+    classifications = torch.stack([item[4] for item in valid_batch])
+    
+    return voxel_data, targets, valid_masks, initial_preds, classifications
+
+# Initialize weights using Xavier/Glorot initialization
+def init_weights(m):
+    if isinstance(m, (nn.Conv3d, nn.Conv1d)):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.Linear):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, (nn.BatchNorm3d, nn.BatchNorm1d)):
+        nn.init.ones_(m.weight)
+        nn.init.zeros_(m.bias)
+
+def train_3dcnn(dataset_dir: str, model_save_path: str, epochs: int = 100, batch_size: int = 16, lr: float = 0.001, 
+                voxel_size: int = 32, score_weight: float = 0.1, class_weight: float = 0.5):
+    """
+    Train the Fast3DCNN model on saved patches.
+    
+    Args:
+        dataset_dir: Directory containing saved patch files
+        model_save_path: Path to save the trained model
+        epochs: Number of training epochs
+        batch_size: Training batch size (reduced due to memory requirements of 3D conv)
+        lr: Learning rate
+        voxel_size: Size of voxel grid
+        score_weight: Weight for the distance prediction loss
+        class_weight: Weight for the classification loss
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Training on device: {device}")
+    
+    # Create dataset and dataloader
+    dataset = VoxelPatchDataset(dataset_dir, voxel_size=voxel_size, augment=True)
+    print(f"Dataset loaded with {len(dataset)} samples")
+    
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4, 
+                           collate_fn=collate_fn, drop_last=True)
+    
+    # Initialize model with score and classification prediction
+    model = Fast3DCNN(input_channels=7, output_dim=3, voxel_size=voxel_size, 
+                      predict_score=True, predict_class=True, num_classes=1)
+    
+    model.apply(init_weights)
+    model.to(device)
+    
+    # Loss functions
+    position_criterion = nn.MSELoss()
+    score_criterion = nn.MSELoss()
+    classification_criterion = nn.BCEWithLogitsLoss()
+    
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5)
+    
+    # Training loop
+    model.train()
+    for epoch in range(epochs):
+        total_loss = 0.0
+        total_pos_loss = 0.0
+        total_score_loss = 0.0
+        total_class_loss = 0.0
+        num_batches = 0
+        
+        for batch_idx, batch_data in enumerate(dataloader):
+            if batch_data is None:  # Skip invalid batches
+                continue
+                
+            voxel_data, targets, valid_masks, initial_preds, classifications = batch_data
+            voxel_data = voxel_data.to(device)  # (batch_size, 7, voxel_size, voxel_size, voxel_size)
+            targets = targets.to(device)  # (batch_size, 3)
+            classifications = classifications.to(device)  # (batch_size,)
+            
+            # Forward pass
+            optimizer.zero_grad()
+            predictions, predicted_scores, predicted_classes = model(voxel_data)
+            
+            # Compute actual distance from predictions to targets
+            actual_distances = torch.norm(predictions - targets, dim=1, keepdim=True)
+            
+            # Only compute position and score losses for samples with GT vertices
+            has_gt_mask = classifications > 0.5
+            
+            if has_gt_mask.sum() > 0:
+                # Position loss only for samples with GT vertices
+                pos_loss = position_criterion(predictions[has_gt_mask], targets[has_gt_mask])
+                score_loss = score_criterion(predicted_scores[has_gt_mask], actual_distances[has_gt_mask])
+            else:
+                pos_loss = torch.tensor(0.0, device=device)
+                score_loss = torch.tensor(0.0, device=device)
+            
+            # Classification loss for all samples
+            class_loss = classification_criterion(predicted_classes.squeeze(), classifications)
+            
+            # Combined loss
+            total_batch_loss = pos_loss + score_weight * score_loss + class_weight * class_loss
+            
+            # Backward pass
+            total_batch_loss.backward()
+            optimizer.step()
+            
+            total_loss += total_batch_loss.item()
+            total_pos_loss += pos_loss.item()
+            total_score_loss += score_loss.item()
+            total_class_loss += class_loss.item()
+            num_batches += 1
+            
+            if batch_idx % 50 == 0:
+                print(f"Epoch {epoch+1}/{epochs}, Batch {batch_idx}, "
+                      f"Total Loss: {total_batch_loss.item():.6f}, "
+                      f"Pos Loss: {pos_loss.item():.6f}, "
+                      f"Score Loss: {score_loss.item():.6f}, "
+                      f"Class Loss: {class_loss.item():.6f}")
+        
+        avg_loss = total_loss / num_batches if num_batches > 0 else 0
+        avg_pos_loss = total_pos_loss / num_batches if num_batches > 0 else 0
+        avg_score_loss = total_score_loss / num_batches if num_batches > 0 else 0
+        avg_class_loss = total_class_loss / num_batches if num_batches > 0 else 0
+        
+        print(f"Epoch {epoch+1}/{epochs} completed, "
+              f"Avg Total Loss: {avg_loss:.6f}, "
+              f"Avg Pos Loss: {avg_pos_loss:.6f}, "
+              f"Avg Score Loss: {avg_score_loss:.6f}, "
+              f"Avg Class Loss: {avg_class_loss:.6f}")
+        
+        scheduler.step()
+
+        # Save model checkpoint every epoch
+        checkpoint_path = model_save_path.replace('.pth', f'_epoch_{epoch+1}.pth')
+        torch.save({
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'epoch': epoch + 1,
+            'loss': avg_loss,
+        }, checkpoint_path)
+    
+    # Save the trained model
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'epoch': epochs,
+    }, model_save_path)
+    
+    print(f"Model saved to {model_save_path}")
+    return model
+
+def load_3dcnn_model(model_path: str, device: torch.device = None, voxel_size: int = 32, predict_score: bool = True) -> Fast3DCNN:
+    """
+    Load a trained Fast3DCNN model.
+    
+    Args:
+        model_path: Path to the saved model
+        device: Device to load the model on
+        voxel_size: Size of voxel grid
+        predict_score: Whether the model predicts scores
+        
+    Returns:
+        Loaded Fast3DCNN model
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    model = Fast3DCNN(input_channels=7, output_dim=3, voxel_size=voxel_size, predict_score=predict_score)
+    
+    checkpoint = torch.load(model_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    
+    model.to(device)
+    model.eval()
+    
+    return model
+
+def predict_vertex_from_patch(model: Fast3DCNN, patch: np.ndarray, device: torch.device = None, voxel_size: int = 32) -> Tuple[np.ndarray, float, float]:
+    """
+    Predict 3D vertex coordinates, confidence score, and classification from a patch using trained 3D CNN.
+    
+    Args:
+        model: Trained Fast3DCNN model
+        patch: Dictionary containing patch data with 'patch_7d' and 'cluster_center' keys
+        device: Device to run prediction on
+        voxel_size: Size of voxel grid
+        
+    Returns:
+        tuple of (predicted_coordinates, confidence_score, classification_score)
+            predicted_coordinates: (3,) numpy array of predicted 3D coordinates
+            confidence_score: float representing predicted distance to GT (lower is better)
+            classification_score: float representing probability of GT vertex presence (0-1)
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    patch_7d = patch['patch_7d']  # (N, 7)
+    
+    # Voxelize the patch
+    voxel_data = voxelize_patch(patch_7d, voxel_size)
+    
+    # Convert to tensor
+    voxel_tensor = torch.from_numpy(voxel_data).float().unsqueeze(0)  # (1, 7, voxel_size, voxel_size, voxel_size)
+    voxel_tensor = voxel_tensor.to(device)
+    
+    # Predict
+    with torch.no_grad():
+        outputs = model(voxel_tensor)
+        
+        if model.predict_score and model.predict_class:
+            position, score, classification = outputs
+            position = position.cpu().numpy().squeeze()
+            score = score.cpu().numpy().squeeze()
+            classification = torch.sigmoid(classification).cpu().numpy().squeeze()
+        elif model.predict_score:
+            position, score = outputs
+            position = position.cpu().numpy().squeeze()
+            score = score.cpu().numpy().squeeze()
+            classification = None
+        elif model.predict_class:
+            position, classification = outputs
+            position = position.cpu().numpy().squeeze()
+            score = None
+            classification = torch.sigmoid(classification).cpu().numpy().squeeze()
+        else:
+            position = outputs
+            position = position.cpu().numpy().squeeze()
+            score = None
+            classification = None
+
+        # Apply offset correction
+        offset = patch['cluster_center']
+        position += offset
+
+        return position, score, classification

hoho_gpu_voxel.batch

@@ -0,0 +1,19 @@
+#!/bin/bash
+#SBATCH --nodes=1 # 1 node
+#SBATCH --ntasks-per-node=1 # 1 tasks per node
+#SBATCH --cpus-per-task=16 # 6 CPUS per task = 12 CPUS per node
+#SBATCH --mem-per-cpu=10G # 8GB per CPU = 96GB per node
+#SBATCH --time=24:00:00 # time limits: 1 hour
+#SBATCH --error=hoho_gpu.err # standard error file
+#SBATCH --output=hoho_gpu.out # standard output file
+#SBATCH --partition=amdgpu # partition name
+#SBATCH --mail-user=skvrnjan@fel.cvut.cz # where send info about job
+#SBATCH --mail-type=ALL # what to send, valid type values are NONE, BEGIN, END, FAIL, REQUEUE, ALL
+#SBATCH --gres=gpu:1
+
+cd /mnt/personal/skvrnjan/hoho/
+module purge
+module load Python/3.10.8-GCCcore-12.2.0
+module load CUDA/12.6.0
+source /mnt/personal/skvrnjan/venvs/hoho/bin/activate
+python train_voxel_cluster.py

train_voxel.py

@@ -0,0 +1,13 @@
+from fast_voxel import train_3dcnn
+import os
+
+if __name__ == "__main__":
+
+    # Load the dataset
+    dataset_path = "/home/skvrnjan/personal/hohocustom/"
+    model_save_path = "/home/skvrnjan/personal/hoho_voxel/"
+
+    os.makedirs(model_save_path, exist_ok=True)
+
+    # Train the model
+    train_3dcnn(dataset_path, model_save_path, epochs=100, batch_size=16, lr=0.001, voxel_size=32, score_weight=0.5, class_weight=0.5)

train_voxel_cluster.py

@@ -0,0 +1,13 @@
+from fast_voxel import train_3dcnn
+import os
+
+if __name__ == "__main__":
+
+    # Load the dataset
+    dataset_path = "/mnt/personal/skvrnjan/hohocustom/"
+    model_save_path = "/mnt/personal/skvrnjan/hoho_voxel/initial.pth"
+
+    os.makedirs(model_save_path, exist_ok=True)
+
+    # Train the model
+    train_3dcnn(dataset_path, model_save_path, epochs=100, batch_size=128, lr=0.001, voxel_size=32, score_weight=0.5, class_weight=0.5)