tpu

model_training_nnn_tpu/data_augmentations.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from scipy.ndimage import gaussian_filter1d
+
+def gauss_smooth(inputs, device, smooth_kernel_std=2, smooth_kernel_size=100,  padding='same'):
+    """
+    Applies a 1D Gaussian smoothing operation with PyTorch to smooth the data along the time axis.
+    Args:
+        inputs (tensor : B x T x N): A 3D tensor with batch size B, time steps T, and number of features N.
+                                     Assumed to already be on the correct device (e.g., GPU).
+        kernelSD (float): Standard deviation of the Gaussian smoothing kernel.
+        padding (str): Padding mode, either 'same' or 'valid'.
+        device (str): Device to use for computation (e.g., 'cuda' or 'cpu').
+    Returns:
+        smoothed (tensor : B x T x N): A smoothed 3D tensor with batch size B, time steps T, and number of features N.
+    """
+    # Get Gaussian kernel
+    inp = np.zeros(smooth_kernel_size, dtype=np.float32)
+    inp[smooth_kernel_size // 2] = 1
+    gaussKernel = gaussian_filter1d(inp, smooth_kernel_std)
+    validIdx = np.argwhere(gaussKernel > 0.01)
+    gaussKernel = gaussKernel[validIdx]
+    gaussKernel = np.squeeze(gaussKernel / np.sum(gaussKernel))
+
+    # Convert to tensor
+    gaussKernel = torch.tensor(gaussKernel, dtype=torch.float32, device=device)
+    gaussKernel = gaussKernel.view(1, 1, -1)  # [1, 1, kernel_size]
+
+    # Prepare convolution
+    B, T, C = inputs.shape
+    inputs = inputs.permute(0, 2, 1)  # [B, C, T]
+    gaussKernel = gaussKernel.repeat(C, 1, 1)  # [C, 1, kernel_size]
+
+    # Perform convolution
+    smoothed = F.conv1d(inputs, gaussKernel, padding=padding, groups=C)
+    return smoothed.permute(0, 2, 1)  # [B, T, C]