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):
  • model_output: Tensor
  • target: Union[Tensor, Tuple, List, Dict] which may contain int, float, str, tuple, list or dict
• Returns one of the following:
  • scalar Tensor
  • tuple or a list containing scalar Tensor (order does not matter); the summation of all items is the total loss
  • dict with str as keys (loss name) and scalar Tensor objects as values. The summation of all dict values is the total loss

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:
  • single Tensor
  • tuple containing Tensor objects
  • list containing Tensor objects
• Second item (the labels) can be any type

(See examples #1 and #2)

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:

• world_size: Optional[int]
• global_rank: Optional[int]
• local_rank: Optional[int]

(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: Tensor, target: Tensor) and returns one of the Metric Return Types

(See example #1 below)

---

• A class implementing (or inheriting from) a TorchMetrics-like API with the following methods:

  • update takes two inputs, (model_output: Tensor, target: Tensor) and updates the state of the object.
    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 a TorchMetrics 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 a forward method which takes two inputs, (model_output: Tensor, target: Tensor) and returns one of the Metric Return Types

(See examples #4 and #5 below)

---

• A stateful class inheriting from the torch.nn.Module that implements the update, compute and reset methods in a TorchMetrics-like API fashion
  • update takes two inputs, (model_output: Tensor, target: Tensor) and updates the state of the object.
    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 scalar Tensor
• tuple or list containing int/float or a scalar Tensor
• dict with str as keys (metric name) and int/float or a scalar Tensor 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()