Nebuly
Search…
Examples and tutorials

A use case for nebulgym

Here we show an example of how you can easily use nebulgym class decorations. To achieve awesome training speed, you can simply add nebulgym decorators (@accelerate_model and @accelerate_dataset) before defining your AI model and dataset.
from typing import List, Callable
​
import torch
from torch.utils.data import Dataset
​
from nebulgym.decorators.torch_decorators import accelerate_model, accelerate_dataset
​
# Add nebulgym class decorator before defining your model.
# This model takes as input an image of resolution 224x224
@accelerate_model()
class CustomModel(torch.nn.Module):
def __init__(self):
super().__init__()
self._avg_pool = torch.nn.AvgPool2d(4)
self._linear = torch.nn.Linear(3136, 1024)
self._relu = torch.nn.ReLU()
self._linears = torch.nn.Sequential(
torch.nn.BatchNorm1d(1024),
torch.nn.Linear(1024, 2048),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(2048),
torch.nn.Linear(2048, 1024),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(1024),
torch.nn.Linear(1024, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512),
torch.nn.Linear(512, 256),
torch.nn.ReLU(),
torch.nn.Linear(256, 2),
)
​
def forward(self, x):
x = self._avg_pool(x).mean(dim=-3).view(-1, 3136)
x = self._relu(self._linear(x))
return self._linears(x)
​
​
# Add nebulgym class decorator before defining your dataset.
@accelerate_dataset()
class CustomDataset(Dataset):
def __init__(self, img_paths: List[str], labelling_func: Callable, reading_func: Callable):
self._images = img_paths
self._labelling_func = labelling_func
self._reading_func = reading_func
​
def __getitem__(self, item):
img_path = self._images[item]
label = self._labelling_func(img_path)
input_tensor = self._reading_func(img_path)
return input_tensor, label
​
def __len__(self):
return len(self._images)
And that's it. Now, as soon as you perform a training run, nebulgym will optimize the full training computing stack, from efficient data loading, to faster forward and backward passes, to earlier convergence.
Export as PDF
Copy link