After the dataloaders are defined, we can iterate them to extract batches of images and their corresponding labels.
This is useful for several purposes, such as visualization, verifying tensor shapes, and more. For example, visualization:
```python
from matplotlib import pyplot as plt
def show(images, labels, classes, rows=6, columns=5):
fig = plt.figure(figsize=(10, 10))
for i in range(1, columns * rows + 1):
fig.add_subplot(rows, columns, i)
plt.imshow(images[i-1].permute(1, 2, 0).clamp(0, 1))
plt.xticks([])
plt.yticks([])
plt.title(f"{classes[labels[i-1]]}")
plt.show()
images_train, labels_train = next(iter(train_dataloader))
show(images_train, labels_train, classes=train_dataloader.dataset.classes)
```
Output:
As can be seen, the images are normalized. The normalization process is defined, among other things, as part of the
default training recipe SuperGradients uses for the CIFAR10 dataset. As we will see in the following section,
SuperGradients makes it a trivial task to override all, or part, of the different dataset and dataloader
parameters, allowing for control over the flexibility vs. ease-of-use tradeoff.
For completion of this section, let's print the tensors' shapes:
```python
print(f'Training image tensor shape: {images_train.shape}')
print(f'Training labels tensor shape: {labels_train.shape}')
```
output:
```
Training image tensor shape: torch.Size([256, 3, 32, 32])
Training labels tensor shape: torch.Size([256])
```
As we can see, the default batch size for the training dataloader is 256.
### 2.B. Override parameters in dataset and dataloaders creation
To showcase the flexibility SuperGradients allows in customizing the different trainer
components, we will override the list of transforms that are used in the dataset. To define a list of
transformations to apply, we will use `torchvision`'s transforms. This also serves to show
the seamless integration SuperGradients allows with different PyTorch components. For the sake of
visualization, the only transform we apply is `ToTensor()`, which simply converts the input
images into PyTorch tensors.
```python
from torchvision import transforms as T
transforms_list = [T.ToTensor()]
vis_dataloader = dataloaders.get("cifar10_train",
dataset_params={"transforms": transforms_list},
dataloader_params={"num_workers": 2})
images, labels = next(iter(vis_dataloader))
show(images, labels, classes=train_dataloader.dataset.classes)
```
Notice that the only difference in the dataloader's definition is that here, the `dataset_params` parameter
is passed as a dictionary defining the parameters to override. The result of running the above code:
The effect of changing the transforms can be seen in the images - now, without normalization, the objects are more
observable in the images.
## 3. Architecture definition
In this example, we train the model with the [ResNet18](https://arxiv.org/abs/1512.03385) architecture.
SuperGradients provides out-of-the-box implementations of many architectures for classification tasks. With just one
line of code we can define a model with the chosen architecture. A list of all available architectures can be found
[here](https://github.com/Deci-AI/super-gradients).
```python
from super_gradients.training import models
from super_gradients.common.object_names import Models
model = models.get(model_name=Models.RESNET18, num_classes=10)
```
Notice that, similar to obtaining a pre-defined dataloader, here we use `super_gradients.training.models`'s `get()`
function. In the above code, two parameters are passed to the function:
* `model_name` - A string defining the model's architecture name, out of the list of architectures SuperGradients
provides.
* `num_classes` - An integer representing the number of classes the model should learn to predict. Affects the
architecture's structure.
Some additional parameters the `get()` function supports:
* `arch_params` - A dictionary used to override the default architecture parameters, such as the number of residual
blocks.
* `checkpoint_path` - A string defining the path to an external checkpoint to be loaded. Can be absolute or relative.
If provided, will automatically attempt to load the checkpoint.
* `pretrained_weights` - A string defining the name of a dataset on which the model was pre-trained on, for
fine-tuning and transfer learning. The `pretrained_weights` and `checkpoint_path` parameters
are mutually exclusive.
For more available parameters, refer to the function's docstring.
In this example, we have defined the model with one of SuperGradient's readily available architectures. As was
already noted in previous sections, SuperGradients is highly compatible with PyTorch. Defining the model's architecture
is not an exception - we can seamlessly use a custom architecture, i.e. a `torch.nn.Module` object, for maximum
flexibility, although it is out of the scope of this example.
## 4. Training setup
We have defined the trainer, datasets, dataloaders, and model architecture. Before we can start training, we need to
define the training parameters. As with the other parameters, SuperGradients provides training parameters optimized for
this use-case. For more recommended training parameters you can have a look at our recipes
[here](https://github.com/Deci-AI/super-gradients/tree/master/src/super_gradients/recipes).
Obtaining the training parameters is as easy as writing a single line of code:
```python
from super_gradients.training import training_hyperparams
training_params = training_hyperparams.get(config_name="training_hyperparams/cifar10_resnet_train_params")
```
We notice the repeatability in the code usage - to obtain the training parameters, we again call
the `get()` function. This function accepts two parameters:
* `config_name` - A string defining the .yaml config filename in the recipes' directory.
* `overriding_params` - An optional parameter, a dictionary used to override the loaded training parameters.
We can print the training parameters to see the different options:
```python
pprint.pprint("Training parameters:")
pprint.pprint(training_params)
```
Output (Training parameters):
```
/* Detailed source-code truncated for AI context efficiency. */
```
As can be seen in the above output, there are numerous options to modify the training parameters
to affect the training process.
It is also possible to change training parameters after obtaining them, for example:
```python
training_params["max_epochs"] = 15
training_params["sg_logger_params"]["launch_tensorboard"] = True
```
## 5. Training, checkpointing, and transfer learning
### 5.A. Training the model
We are all set to start training our model. Simply plug in the model, training and validation dataloaders,
and training parameters into the trainer's `train()` function:
```python
trainer.train(model=model,
training_params=training_params,
train_loader=train_dataloader,
valid_loader=valid_dataloader)
```
The training progress will be printed to the screen:
```
[2023-02-01 20:57:27] INFO - sg_trainer_utils.py - TRAINING PARAMETERS:
- Mode: Single GPU
- Number of GPUs: 1 (4 available on the machine)
- Dataset size: 50000 (len(train_set))
- Batch size per GPU: 256 (batch_size)
- Batch Accumulate: 1 (batch_accumulate)
- Total batch size: 256 (num_gpus * batch_size)
- Effective Batch size: 256 (num_gpus * batch_size * batch_accumulate)
- Iterations per epoch: 195 (len(train_set) / total_batch_size)
- Gradient updates per epoch: 195 (len(train_set) / effective_batch_size)
[2023-02-01 20:57:27] INFO - sg_trainer.py - Started training for 15 epochs (0/14)
Train epoch 0: 100%|██████████| 196/196 [00:18<00:00, 10.51it/s, Accuracy=0.262, CrossEntropyLoss=2.37, Top5=0.787, gpu_mem=0.371]
Validation epoch 0: 100%|██████████| 20/20 [00:03<00:00, 6.12it/s]
===========================================================
SUMMARY OF EPOCH 0
├── Training
│ ├── Accuracy = 0.262
│ ├── CrossEntropyLoss = 2.3702
│ └── Top5 = 0.787
└── Validation
├── Accuracy = 0.3459
├── CrossEntropyLoss = 1.8811
└── Top5 = 0.871
===========================================================
```
At the beginning of the training, a summary of the training parameters is printed, where we can see
the training mode (CPU/single GPU/distributed training), the number of GPUs used, the training dataset
size, and more. The progress of each epoch's training and validation is displayed, along with the tracked metrics
(defined as part of the training recipe): accuracy, loss value, top5 error, and GPU memory consumption.
At the end of each epoch, a summary of the training and validation metrics is displayed, and in later epochs,
a comparison with the previous epochs is provided:
```
===========================================================
SUMMARY OF EPOCH 15
├── Training
│ ├── Accuracy = 0.7594
│ │ ├── Best until now = 0.7458 (↗ 0.0136)
│ │ └── Epoch N-1 = 0.7458 (↗ 0.0136)
│ ├── CrossEntropyLoss = 0.686
│ │ ├── Best until now = 0.7187 (↘ -0.0327)
│ │ └── Epoch N-1 = 0.7187 (↘ -0.0327)
│ └── Top5 = 0.9867
│ ├── Best until now = 0.9849 (↗ 0.0019)
│ └── Epoch N-1 = 0.9849 (↗ 0.0019)
└── Validation
├── Accuracy = 0.7425
│ ├── Best until now = 0.746 (↘ -0.0035)
│ └── Epoch N-1 = 0.7306 (↗ 0.0119)
├── CrossEntropyLoss = 0.7331
│ ├── Best until now = 0.7315 (↗ 0.0016)
│ └── Epoch N-1 = 0.8048 (↘ -0.0717)
└── Top5 = 0.9831
├── Best until now = 0.9838 (↘ -0.0007)
└── Epoch N-1 = 0.9818 (↗ 0.0013)
===========================================================
```
At the end of each epoch, the different logs and checkpoints are saved in the path defined by `ckpt_root_dir` and
`experiment_name`.
In the epoch summary shown above, we can see that the validation accuracy is 73%, which is
not very high. To get better insights as to what is happening, we turn to the tensorboard logs.
### 5.B. Tensorboard logs
To view the experiment's tensorboard logs, type the following command in the terminal from the
experiment's path:
```bash
tensorboard --logdir='.'
```
(Alternatively, run the command from anywhere with the experiment's full path).
SuperGradients logs many useful metrics to tensorboard, including CPU and GPU usage, learning rate
scheduling, training and validation losses and other metrics, and many more. For the purpose of this example,
let us examine the training and validation loss:
As can be seen in the graphs, the training (and validation) loss did not converge before training ended. This means
that training the model for additional epochs will probably improve its performance. Earlier, when modifying the
training parameters, we set `max_epochs = 15`. Let us continue training the model for an additional 10 epochs.
### 5.C. Continue training from a checkpoint
To continue training from a checkpoint, we utilize the `models.get()` function's `checkpoint_path` parameter.
The provided checkpoint path includes the checkpoint file we wish to load. In this example, since we want
to continue from the last checkpoint, we will load the `ckpt_latest.pth` checkpoint. Additionally, we want to let
the trainer know that we are continuing training and not starting from the first epoch. This is done
by setting the `resume` training parameter to `True`. Finally, we set the new `max_epochs` training
parameter, and train the model once more.
```python
import os
model = models.get(model_name=Models.RESNET18,
num_classes=10,
checkpoint_path=os.path.join(CHECKPOINT_DIR, experiment_name, 'ckpt_latest.pth'))
training_params["resume"] = True
training_params["max_epochs"] = 25
trainer.train(model=model,
training_params=training_params,
train_loader=train_dataloader,
valid_loader=valid_dataloader)
```
We can see that the training continues for 10 epochs, resuming from epoch 15:
```
[2023-02-01 21:21:16] INFO - sg_trainer_utils.py - TRAINING PARAMETERS:
- Mode: Single GPU
- Number of GPUs: 1 (4 available on the machine)
- Dataset size: 50000 (len(train_set))
- Batch size per GPU: 256 (batch_size)
- Batch Accumulate: 1 (batch_accumulate)
- Total batch size: 256 (num_gpus * batch_size)
- Effective Batch size: 256 (num_gpus * batch_size * batch_accumulate)
- Iterations per epoch: 195 (len(train_set) / total_batch_size)
- Gradient updates per epoch: 195 (len(train_set) / effective_batch_size)
[2023-02-01 21:21:16] INFO - sg_trainer.py - Started training for 10 epochs (15/24)
Train epoch 15: 100%|██████████| 196/196 [00:18<00:00, 10.52it/s, Accuracy=0.764, CrossEntropyLoss=0.668, Top5=0.987, gpu_mem=0.422]
Validation epoch 15: 100%|██████████| 20/20 [00:03<00:00, 6.10it/s]
===========================================================
SUMMARY OF EPOCH 15
├── Training
│ ├── Accuracy = 0.7644
│ ├── CrossEntropyLoss = 0.6684
│ └── Top5 = 0.9865
└── Validation
├── Accuracy = 0.7539
├── CrossEntropyLoss = 0.7271
└── Top5 = 0.9841
===========================================================
```
Finally, the model stops training after completing 25 epochs:
```
===========================================================
SUMMARY OF EPOCH 25
├── Training
│ ├── Accuracy = 0.8177
│ │ ├── Best until now = 0.8147 (↗ 0.003)
│ │ └── Epoch N-1 = 0.8147 (↗ 0.003)
│ ├── CrossEntropyLoss = 0.5211
│ │ ├── Best until now = 0.5281 (↘ -0.007)
│ │ └── Epoch N-1 = 0.5281 (↘ -0.007)
│ └── Top5 = 0.9921
│ ├── Best until now = 0.9919 (↗ 0.0002)
│ └── Epoch N-1 = 0.9919 (↗ 0.0002)
└── Validation
├── Accuracy = 0.8201
│ ├── Best until now = 0.7873 (↗ 0.0328)
│ └── Epoch N-1 = 0.7534 (↗ 0.0667)
├── CrossEntropyLoss = 0.525
│ ├── Best until now = 0.6145 (↘ -0.0895)
│ └── Epoch N-1 = 0.7517 (↘ -0.2266)
└── Top5 = 0.9907
├── Best until now = 0.9883 (↗ 0.0024)
└── Epoch N-1 = 0.983 (↗ 0.0077)
===========================================================
```
We can see that the validation accuracy is now 82%. Much better.
### 5.C. Transfer learning
So far, we trained a model from scratch. More formally, the model's weights were randomly initialized. For easy tasks
and large datasets this is usually sufficient. In other cases, especially when not a lot of data is available for
training, we would like to take advantage of knowledge gained from other sources. This is called transfer learning,
and it has many forms and variations. In this example we will take a look at the simplest form of transfer learning:
fine-tuning a model initialized with pre-trained weights. Specifically, we will initialize our model
with weights pre-trained on the [ImageNet dataset](https://www.image-net.org/).
SuperGradients provides a variety of pre-trained weights readily available for fine-tuning different models. To
initialize our model with pre-trained weights provided by SuperGradients, only a small change to the existing code
is needed:
```
model = models.get(model_name=Models.RESNET18, num_classes=10, pretrained_weights="imagenet")
```
In the above code, we provided the `pretrained_weights` parameter to the `models.get()` function. This parameter accepts
a string representing the name of the dataset that the weights were pre-trained on. Note that this
parameter is mutually exclusive with the `checkpoint_path` parameter.
The rest of the training pipeline is the same as above. For comparison with the previous model, we will train this model
for 25 epochs also. Let us look how the model performed:
```
===========================================================
SUMMARY OF EPOCH 25
├── Training
│ ├── Accuracy = 0.8242
│ │ ├── Best until now = 0.8267 (↘ -0.0025)
│ │ └── Epoch N-1 = 0.8267 (↘ -0.0025)
│ ├── CrossEntropyLoss = 0.5035
│ │ ├── Best until now = 0.4998 (↗ 0.0037)
│ │ └── Epoch N-1 = 0.4998 (↗ 0.0037)
│ └── Top5 = 0.9924
│ ├── Best until now = 0.9917 (↗ 0.0007)
│ └── Epoch N-1 = 0.9917 (↗ 0.0007)
└── Validation
├── Accuracy = 0.8377
│ ├── Best until now = 0.8062 (↗ 0.0315)
│ └── Epoch N-1 = 0.806 (↗ 0.0317)
├── CrossEntropyLoss = 0.4834
│ ├── Best until now = 0.5731 (↘ -0.0897)
│ └── Epoch N-1 = 0.5785 (↘ -0.0952)
└── Top5 = 0.9924
├── Best until now = 0.9903 (↗ 0.0021)
└── Epoch N-1 = 0.9903 (↗ 0.0021)
===========================================================
```
As we can see, the validation accuracy improved by 1.7% compared to the randomly initialized model.
To achieve a greater improvement with pre-trained weights, sometimes careful tuning of the training hyperparameters
is required.
## 6. Predictions with the trained model
Now that we have a trained model with reasonable performance, we can use it to make predictions on new data.
First, let's import some packages:
```python
from PIL import Image
import torch
import numpy as np
import requests
```
Next, we load the model with the trained weights, and put it into evaluation mode. Notice that now we load the
`ckpt_best.pth` checkpoint.
```python
model = models.get(model_name=Models.RESNET18,
num_classes=10,
checkpoint_path=os.path.join(CHECKPOINT_DIR, experiment_name, 'ckpt_best.pth'))
model.eval()
```
We want to test the model on an image of one the classes the model was trained on. As an example, let us test how the
model handles an image of a frog, taken from the [Aquarium of the Pacific website](https://www.aquariumofpacific.org/).
The loaded image must undergo the same transformations as the training images for the model to work well:
```python
url = "https://www.aquariumofpacific.org/images/exhibits/Magnificent_Tree_Frog_900.jpg"
image = np.array(Image.open(requests.get(url, stream=True).raw))
transforms = T.Compose([
T.ToTensor(),
T.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)),
T.Resize((32, 32))
])
input_tensor = transforms(image).unsqueeze(0).to(next(model.parameters()).device)
```
Next, to obtain the model's predictions we simply run the following line of code:
```python
predictions = model(input_tensor)
```
Let's see what the model predicted:
```python
plt.xlabel(train_dataloader.dataset.classes[torch.argmax(predictions)])
plt.imshow(image)
plt.show()
```
As we can see, the model correctly predicted that the input image is an image of a frog.
## 7. Complete code
For completion of this example, we provide a complete working code for training, continuing training from
a saved checkpoint, and predicting with the trained model. Simply change the `CHECKPOINT_DIR` variable and run the
script:
```
/* Detailed source-code truncated for AI context efficiency. */
```
---
### Documentation/Source/Example Training An External Model (documentation/source/Example_Training-an-external-model.md)
# Training an external model
In this example we will use SuperGradients to train a deep learning segmentation model to extract human portraits from
images, i.e., to remove the background from the image.
We will show how SuperGradients allows seamless integration of
an external model, dataset, loss function, and metric into the training pipeline.
## Quick installation
For this example, the only necessary package is super-gradients. Installing super-gradients will also
install all dependencies required to run the code in this example.
```bash
pip install super-gradients
```
## 1. Dataset
The dataset we will use in this example is the
[AISegment dataset](https://github.com/aisegmentcn/matting_human_datasets), available to download for free from
[Kaggle](https://www.kaggle.com/datasets/laurentmih/aisegmentcom-matting-human-datasets).
The dataset contains 34,427 RGB images of human portraits and their corresponding **soft masks**.
The provided images are center-cropped to a unified shape of 600x800.
### 1.A. Data preparation
The original structure of the data on the disk is as follows:
```
data
└───aisegment-matting
│ └───matting
│ │ └───1803290511
│ │ │ └───matting_00000000
│ │ │ │ 1803290511-00000459.png
│ │ │ │ 1803290511-00000458.png
│ │ │ │ ..
│ │ └───1803290444
│ │ ..
│ └───clip_img
│ │ └───1803290511
│ │ │ └───clip_00000000
│ │ │ │ 1803290511-00000459.jpg
│ │ │ │ 1803290511-00000458.jpg
│ │ │ │ ..
│ │ └───1803290444
│ │ ..
```
In the structure shown above, the `clip_img` directory contains the portrait images, divided into many
sub-directories. The `matting` directory is structured similar to the `clip_img` directory, containing
the corresponding masks. This structure is not particularly convenient for data loading, therefore
we will first rearrange the data with the following script, which only requires changing the `data_path` and
`out_path` variables:
```
/* Detailed source-code truncated for AI context efficiency. */
```
The original data may have images without their corresponding masks. Therefore, the above code first finds all existing
image-mask pairs. It then splits the data into train/validation/test sets, where the size of each of the validation and
test sets is 10% of the total number of samples. The samples in each of the sets are then copied to the output path,
structured as follows:
```
data
└───train
│ 1803290511-00000029.jpg
│ 1803290511-00000029.png
│ ..
└───val
│ 1803290444-00000292.jpg
│ 1803290444-00000292.png
│ ..
└───test
│ 1803290443-00000294.jpg
│ 1803290443-00000294.png
│ ..
```
Each of the train/val/test directories contain all image (*.jpg) files and their corresponding mask (*.png) files.
### 1.B. PyTorch Dataset
In some cases, we may want to have full control over the process of loading and pre-processing the training data.
SuperGradients is fully compatible with PyTorch data loaders, which allows for seamless integration of custom dataset
implementations for maximum flexibility.
We will first present the complete Dataset class implementation, and then break it down to fully understand what is
going on.
```
/* Detailed source-code truncated for AI context efficiency. */
```
First, we can see that our dataset class inherits from the `torch.utils.data.Dataset` class. To initialize the dataset
object, two parameters must be provided:
* `data_path` - the full path to the data's root directory
* `split` - a string indicating whether this is the 'train', 'val', or 'test' split of the data
Additional parameters include `input_height` and `input_width`, the fixed height and width, respectively, for resizing
the input images and masks. These are set to 256 by default. At the end of the initialization function, the image and
mask path pairs are determined according to the split.
Next, let's see what happens when we retrieve an item from the dataset via the `__getitem__()` function.
First, an image and its corresponding mask are loaded according to the `idx` parameter:
```python
image = Image.open(self.path_pairs[idx][0]).convert('RGB')
mask = Image.open(self.path_pairs[idx][1]).split()[-1]
```
Notice that we take only the mask's last channel. The mask's original color format is RGBA, and the alpha channel should be
used as the soft mask for segmentation according to the
[dataset's description](https://www.kaggle.com/datasets/laurentmih/aisegmentcom-matting-human-datasets?resource=download).
Next, we apply transformations to the image and mask according to the split:
```python
if self.split == 'train':
seg_transforms = torch_transforms.Compose([
SegRandomFlip(prob=0.5),
SegColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
SegResize(h=self.input_height, w=self.input_width)
])
else:
seg_transforms = SegResize(h=self.input_height, w=self.input_width)
transformed_pair = seg_transforms({"image": image, "mask": mask})
image, mask = transformed_pair['image'], transformed_pair['mask']
```
On training images and masks we apply data augmentation: a random horizontal flip with probability 0.5, and
random color jitter which randomly changes the image's brightness, contrast, and saturation. Both the training and the
other splits' images and masks undergo resizing according to `input_height` and `input_width`.
The transforms in the above code are all SuperGradients transforms, which are built upon PyTorch's
`torchvision.transforms`. This allows for maximum compatibility with PyTorch components. For example, we can see that
the transforms are sequentially composed using torchvision's `Compose` class. Notice also that the transform names are
all prefixed with `Seg` - meaning that they are specifically designed for segmentation data. These transforms take care
to apply the same transformation to the image and the mask when required, or only apply the transformation to the image
otherwise. For example, if `SegRandomFlip()` flips the image, the mask will be flipped as well. `SegColorJitter()` only
transforms the image. `SegResize()` resizes both the image and the mask.
Next, we convert the image and the mask to PyTorch Tensors, and normalize the image using the `NORMALIZATION_MEANS` and
`NORMALIZATION_STDS`:
```python
image_transform = torch_transforms.Compose([
torch_transforms.ToTensor(),
torch_transforms.Normalize(self.NORMALIZATION_MEANS, self.NORMALIZATION_STDS)
])
mask_transform = torch_transforms.ToTensor()
image_tensor, mask_tensor = image_transform(image), mask_transform(mask)
```
These transforms are applied regardless of the current dataset split. Notice that here we use torchvision transforms.
This serves to show the high degree of flexibility SuperGradients allows, as its transforms are based on PyTorch's
torchvision transforms which may be used interchangeably.
### 1.C. Data visualization
To conclude this section, let's visualize some images and their masks to test our AISegmentDataset implementation.
First, we instantiate two dataset objects, for the training and validation splits, and extract the first sample from
each:
```python
import matplotlib.pyplot as plt
data_path = '/path/to/arranged/data/dir'
train_dataset = AISegmentDataset(data_path=data_path, split='train')
val_dataset = AISegmentDataset(data_path=data_path, split='val')
train_image, train_mask = train_dataset[0]
val_image, val_mask = val_dataset[0]
```
Let's first visualize the validation image and mask:
```python
figure = plt.figure()
figure.add_subplot(1, 2, 1)
plt.title("Image")
plt.axis("off")
plt.imshow(val_image.permute(1, 2, 0))
figure.add_subplot(1, 2, 2)
plt.title("Mask")
plt.axis("off")
plt.imshow(val_mask.squeeze(), cmap='gray')
plt.show()
```
And the training image and mask:
```python
figure = plt.figure()
figure.add_subplot(1, 2, 1)
plt.title("Image")
plt.axis("off")
plt.imshow(train_image.permute(1, 2, 0))
figure.add_subplot(1, 2, 2)
plt.title("Mask")
plt.axis("off")
plt.imshow(train_mask.squeeze(), cmap='gray')
plt.show()
```
We can see the effect of the color jitter transform on the image.
## 2. Model architecture
For this example we will employ the [U-Net](https://arxiv.org/abs/1505.04597) architecture. U-Net and its variants are a
popular choice for many image segmentation tasks. It is a fully-convolutional architecture,
consisting of an expanding and a contracting path with skip connections between the encoder
and decoder blocks. In this example we will demonstrate how we can easily integrate an external PyTorch model
as part of SuperGradients' training pipeline. To this end, we will use a
[U-Net implementation](https://github.com/mateuszbuda/brain-segmentation-pytorch) loaded directly from
[PyTorch Hub](https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/):
```python
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet', in_channels=3, out_channels=1,
init_features=32, pretrained=False)
```
Since our model's inputs are RGB images, we set `in_channels=3`. This is a binary segmentation task, therefore
`out_channels=1`. The `init_features` parameter determines the number of kernels in the first block's convolution
layers. The number of kernels is doubled in each consecutive encoder block. We also set `pretrained=False` since
in this case we do not want to use pre-trained weights.
Let's check our model's type:
```python
print(type(model))
```
```
We can also check the validation set's IoU metric's value:
## 7. Predictions with the trained model
Now that we have a trained model we can use it to make predictions on the test set. First, let's instantiate a test
dataset:
```python
test_dataset = AISegmentDataset(data_path=data_path, split='test')
```
By instantiating the test dataset, all required pre-processing is already handled for us. Let's choose a single sample
from the dataset:
```python
image, mask = test_dataset[0]
image, mask = image.unsqueeze(0), mask.unsqueeze(0)
```
Notice that we have added the batch dimension to the tensors. Next, we set the model to evaluation mode, and obtain
the predicted mask. Since our model learned to predict soft masks, we apply a threshold of 0.5 to binarize the mask:
```python
model.eval()
pred = model(image)
pred = torch.gt(pred, 0.5).long()
```
Finally, let's apply the predicted mask to the image:
```python
masked_image = image*pred
```
Now let's visualize the image, mask, and masked image to see how our model performed:
```python
figure = plt.figure()
figure.add_subplot(1, 3, 1)
plt.title("Image")
plt.axis("off")
plt.imshow(image.detach().squeeze(0).permute(1, 2, 0))
figure.add_subplot(1, 3, 2)
plt.title("Mask")
plt.axis("off")
plt.imshow(pred.detach().squeeze(0).permute(1, 2, 0), cmap='gray')
figure.add_subplot(1, 3, 3)
plt.title("Masked Image")
plt.axis("off")
plt.imshow(masked_image.detach().squeeze(0).permute(1, 2, 0))
plt.show()
```
---
### Documentation/Source/Experiment Management (documentation/source/experiment_management.md)
# Experiment Management
## Outline
1. [Core Concepts](#core-concepts)
- [Checkpoint Root Directory](#checkpoint-root-directory-ckpt_root_dir)
- [Experiments](#experiments-experiment_name)
- [Runs](#runs-run_id)
2. [File Structure of Experiments](#file-structure-of-experiments)
3. [Utilities for Experiment Management](#utilities)
- [Get the Absolute Path of a Run Directory](#a-get-the-absolute-path-of-a-run-directory)
- [Retrieve the Latest Run ID](#b-get-the-latest-run-id)
## Core Concepts
### Checkpoint Root Directory (`ckpt_root_dir`)
- The main directory where all experiment outputs are housed.
### Experiments (`experiment_name`)
- Symbolizes a distinct training recipe or configuration.
- Alter the `experiment_name` for transparency when updating your training recipe.
- Each training under the same `experiment_name` has its individual `run` directory, ensuring no overwrites.
### Runs (`run_id`)
- Every individual training session is termed as a `run`.
- A unique `run_id` is generated for every training, regardless of identical parameters.
- Different trainings under the same `experiment_name` maintain distinct logs and checkpoints, courtesy of their separate run directories.
## File Structure of Experiments
```
Predict poses with YoloNAS Pose Model
.yaml configuration file:.yaml configuration file:.yaml configuration file:.yaml configuration file:
