Blealtan / efficient-kan

An efficient pure-PyTorch implementation of Kolmogorov-Arnold Network (KAN).
MIT License
3.49k stars 306 forks source link

CUDA out of memory. #24

Closed ybu-lxd closed 1 month ago

ybu-lxd commented 1 month ago
from src.efficient_kan.kan import KAN
import torch
net = KAN([1152,1152*4,1152]).to("cuda")
x = torch.rand(size=(4096*4,1152)).to("cuda")
net(x)

I found that if the hidden layer is too large, the problem of CUDA out of memory will occur.

Arturossi commented 1 month ago

This happens because your model demands more memory than your GPU can provide. Try to use batches and adjust their size or modify your topology so everything can fit in your memory.

This approach most likely will impact in your network results, each one in its own way.

Example for batching that I have done in one of my projects:

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, features, target):
        self.features = features
        self.target = target

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.target[idx]

batch_size = 32 # Example of batch size

loader = DataLoader(
                dataset = CustomDataset(X_train, y_train), 
                batch_size = batch_size, 
                shuffle = True # Shuffle the entire dataset before each epoch to help preventing overfitting
            )

for i , (features, targets) in enumerate(loader):
    # Process your loader here
Blealtan commented 1 month ago

See #23.