CCO Input Requirements
Loss Function
When writing a loss function or a callable class for CCO, it must satisfy the following rules:
- Takes two input arguments (order matters):
- Returns one of the following:
Examples
from torch import Tensor
from typing import Union, Tuple, List, Dict
def loss_fn(model_output: Tensor, target: Tensor) -> Union[Tensor, Tuple[Tensor, ...],List[Tensor, ...], Dict[str, Tensor]]:
# calculate loss...
loss: Tensor = some_scalar_tensor
# returns a loss as a tensor
return loss
# or a tuple of scalar tensors
return loss_1, loss_2, loss_3
# or a dictionary with scalar tensors as dictionary values
return {"cls_loss": loss_1, "box_loss": loss_2, "obj_loss": loss_3}
Dataset / Dataloader
Dataset / Dataloader output format
When writing a dataset class for CCO, it must return a tuple of length two (order matters)
- First item (the input data to the model) can be one of the following:
- Second item (the labels) can be any type
Dataloader input for Distributed Compression
Distributed Training can be used by setting multi_gpu=True
in the cc.PyTorchCompressionEngine.optimize()
or cc.PyTorchCompressionEngine.resume()
.
If you need finer control over the Dataset in the context of Distributed Compression, you may set the Dataloader initializer function so that it takes the following arguments by either keyword or as positional arguments:
(See example #3)
Example #1
In this example we return a tuple containing two Tensor
objects:
import torch
from torch import Tensor
from typing import Tuple
class CustomDataset(torch.utils.data.Dataset):
...
def __getitem__(self, x) -> Tuple[Tensor, Tensor]:
...
return (input_tensor, label_tensor)
Example #2
If you require auxiliary information to be included, you may set the second item in the returned tuple to be a dictionary or a tuple (or any other type) containing any auxiliary information necessary as shown below:
from torch import Tensor
from typing import Tuple
from torch.utils.data import Dataset
class CustomDataset(Dataset):
...
def __getitem__(self, x) -> Tuple[Tensor, Tensor]:
...
return (input_tensor, label_tensor)
...
Example #3
Setting the dataloader initializer function that can take Distributed Compression related arguments for finer control during the Multi-GPU Distributed Compression:
from torch.utils.data import DataLoader
from typing import Optional
from clika_compression import PyTorchCompressionEngine
from clika_compression.settings import generate_default_settings
...
def get_train_loader(world_size: Optional[int], global_rank: Optional[int],
local_rank: Optional[int], train_dataset):
return DataLoader(train_dataset,batch_size=64, shuffle=True)
def get_eval_loader(world_size: Optional[int], global_rank: Optional[int],
local_rank: Optional[int], eval_dataset):
return DataLoader(eval_dataset,batch_size=64, shuffle=True)
...
engine = PyTorchCompressionEngine()
settings = generate_default_settings() # Can be configured after creation
# Start the "CCO" to compress model to generate a `.pompom` file
engine.optimize( # final is the path to the `.pompom` file from the latest epoch
settings=settings, # the CCO settings
output_path='outputs',
model=model, # the model to compress
model_compile_settings=mcs,
init_training_dataset_fn=get_train_loader, # a function that returns the training dataloader
init_evaluation_dataset_fn=get_eval_loader,
multi_gpu=True) # use Multi-GPU Distributed Compression
Metric Function / Class
When writing a metric function or class for CCO,
it needs to be one of the following objects
(all metrics will be shown in the CCO output log):
- A function taking two inputs -
(model_output:
and returns one of the Metric Return TypesTensor
, target:Tensor
)
(See example #1 below)
A class implementing (or inheriting from) a
TorchMetrics
-like API with the following methods:update
takes two inputs,(model_output:
and updates the state of the object.Tensor
, target:Tensor
)
It may return one of the Metric Return Types (return values will be shown on the output log training steps and epoch summary)compute
takes no inputs, calculate the final metric, and returns it as one of the Metric Return Types (return values will be shown on the output log epoch summary)reset
- when not inheriting from aTorchMetrics
object, this method should be implemented. The method resets the metric object's state attributes to their default values. It is called at the end of each epoch and takes no inputs
(See examples #2 and #3 below)
- A stateless class inheriting from the
torch.nn.Module
that implements aforward
method which takes two inputs,(model_output:
and returns one of the Metric Return TypesTensor
, target:Tensor
)
(See examples #4 and #5 below)
- A stateful class inheriting from the
torch.nn.Module
that implements theupdate
,compute
andreset
methods in aTorchMetrics
-like API fashionupdate
takes two inputs,(model_output:
and updates the state of the object.Tensor
, target:Tensor
)
It may return one of the Metric Return Types (return values will be shown on the output log training steps and epoch summary)compute
takes no inputs, calculate the final metric, and returns it as one of the Metric Return Types (return values will be shown on the output log epoch summary)reset
resets the metric object's state attributes to their default values. It is called at the end of each epoch and takes no inputs.
(See example #6 below)
Metric Return Types
These are types that are the standard for metric method output.
int
/float
or a scalarTensor
tuple
orlist
containingint
/float
or a scalarTensor
dict
withstr
as keys (metric name) andint
/float
or a scalarTensor
as values
Examples
import torch
import torchmetrics
from torch import Tensor
from typing import Dict
# Example 1 - lambda function example that takes two arguments and returns a tensor scalar
metric_fn = lambda model_output, target: torch.mean((torch.argmax(model_output, 1) == target).float()) * 100
# ============================ #
# Example 2 - inheriting from `torchmetrics.classification.MulticlassAccuracy`
class Accuracy(torchmetrics.classification.MulticlassAccuracy):
...
def update(self, model_output: Tensor, target: Tensor) -> None:
super().update(model_output, target)
def compute(self):
result = super().compute()
return result * 100
metric_fn = Accuracy(num_classes=num_classes)
# ============================ #
# Example 3 - a custom class implementing TorchMetrics-like API
class Accuracy(object):
def __init__(self):
self.avgs = []
def update(self, model_output: Tensor, target: Tensor) -> None:
batch_mean = torch.mean((torch.argmax(model_output, 1) == target).float())
self.avgs.append(batch_mean.item())
def compute(self):
return sum(self.avgs) / len(self.avgs)
def reset(self):
# resting the object state attributes to default values
self.avgs = []
metric_fn = Accuracy()
# ============================ #
# Example 4 - using a torch.nn.Module class directly as metric (already satisfies the return inputs and outputs constraints) in a stateless fashion
metric_fn = torch.nn.cNLLLoss()
# ============================ #
# Example 5 - a stateless class inheriting from torch.nn.Module
class Accuracy(torch.nn.Module):
def forward(self, model_output: Tensor, target: Tensor)-> Dict[str, Tensor]:
# calculate metric values
num_correct_predictions = torch.sum(torch.argmax(model_output, 1) == target)
accuracy = torch.mean((torch.argmax(model_output, 1) == target).float())
# return a dict with strings as keys and scalar tensors as values
res = {"num_correct_predictions": num_correct_predictions, "accuracy": accuracy}
return res
metric_fn = Accuracy()
# ============================ #
# Example 6 - a stateful class inheriting from torch.nn.Module
from typing import Union, Tuple, List, Dict
class Accuracy(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.register_buffer("correct", torch.Tensor([0]))
self.register_buffer("total", torch.Tensor([0]))
def compute(self) -> Union[Tensor, Tuple, List, Dict]:
return self.correct / self.total
def reset(self) -> None:
self.correct.copy_(0)
self.total.copy_(0)
def update(self, preds, targets) -> Union[Tensor, Tuple, List, Dict, None]:
# assuming preds is (N, C) and targets is (N,) and categories is (,C)
correct = (preds.argmax(-1) == targets).sum().item()
total = preds.shape[0]
self.correct += correct
self.total += total
return {"batch_acc": correct / total}
metric_fn = Accuracy()