mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
import torch
|
|
import random
|
|
import numpy as np
|
|
|
|
def shuffle_tensor_along_axis(tensor, axis=0, seed=None):
|
|
"""
|
|
Shuffle a tensor along a specified axis without affecting the global random state.
|
|
|
|
Args:
|
|
tensor (torch.Tensor): The input tensor to shuffle
|
|
axis (int, optional): The axis along which to shuffle. Defaults to 0.
|
|
seed (int, optional): Random seed for reproducibility. Defaults to None.
|
|
|
|
Returns:
|
|
torch.Tensor: The shuffled tensor
|
|
"""
|
|
# Clone the tensor to avoid in-place modifications
|
|
shuffled_tensor = tensor.clone()
|
|
|
|
# Store original random states
|
|
torch_state = torch.get_rng_state()
|
|
np_state = np.random.get_state()
|
|
py_state = random.getstate()
|
|
|
|
try:
|
|
# Set seed if provided
|
|
if seed is not None:
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
|
|
# Get the size of the dimension to shuffle
|
|
dim_size = tensor.shape[axis]
|
|
|
|
# Generate random indices for shuffling
|
|
indices = torch.randperm(dim_size)
|
|
|
|
# Create a slice object to shuffle along the specified axis
|
|
slices = [slice(None)] * tensor.dim()
|
|
slices[axis] = indices
|
|
|
|
# Apply the shuffle
|
|
shuffled_tensor = tensor[slices]
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error during shuffling: {e}")
|
|
|
|
finally:
|
|
# Restore original random states
|
|
torch.set_rng_state(torch_state)
|
|
np.random.set_state(np_state)
|
|
random.setstate(py_state)
|
|
|
|
return shuffled_tensor |