Pytorch Generator#

A code generator for PyTorch is also provided by BESSER. PyTorch is an open-source machine learning framework developed by Meta (Facebook) that is widely used for deep learning. Our Generator transforms B-UML Neural Network models into PyTorch code, allowing you to create neural networks based on your B-UML specifications.

To use the PyTorch generator, you need to create a PytorchGenerator object, provide the B-UML Neural Network model, and use the generate method as follows:

from besser.generators.nn.pytorch.pytorch_code_generator import PytorchGenerator

pytorch_model = PytorchGenerator(
    model=nn_model, output_dir="output_folder", generation_type="subclassing"
)
pytorch_model.generate()

The configuration parameters for the PytorchGenerator are as follows:

  • model: The neural network model.

  • output_dir: The name of the output directory where the pytorch_nn_subclassing.py file will be generated.

  • generation_type: The type of NN architecture. Either subclassing or sequential.

The pytorch_nn_subclassing.py file will be generated inside output_folder and it will look as follows:

  1"""PyTorch code generated based on BUML."""
  2import torch
  3from datetime import datetime
  4
  5
  6from torch import nn
  7from torchvision import datasets, transforms
  8
  9from sklearn.metrics import classification_report 
 10
 11
 12# Define the network architecture
 13class NeuralNetwork(nn.Module):
 14    def __init__(self):
 15        super().__init__()
 16        self.l1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=(3, 3), stride=(1, 1), padding=0)
 17        self.actv_func_relu = nn.ReLU()
 18        self.l2 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0)
 19        self.l3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=0)
 20        self.l4 = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0)
 21        self.l5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=0)
 22        self.l6 = nn.Flatten(start_dim=1, end_dim=-1)
 23        self.l7 = nn.Linear(in_features=1024, out_features=64)
 24        self.l8 = nn.Linear(in_features=64, out_features=10)
 25
 26
 27    def forward(self, x):
 28        x = self.l1(x)
 29        x = self.actv_func_relu(x)
 30        x = self.l2(x)
 31        x = self.l3(x)
 32        x = self.actv_func_relu(x)
 33        x = self.l4(x)
 34        x = self.l5(x)
 35        x = self.actv_func_relu(x)
 36        x = self.l6(x)
 37        x = self.l7(x)
 38        x = self.actv_func_relu(x)
 39        x = self.l8(x)
 40        return x
 41
 42
 43# Dataset preparation
 44IMAGE_SIZE = (32, 32)
 45transform = transforms.Compose([
 46    transforms.Resize(IMAGE_SIZE),
 47	transforms.ToTensor()
 48    ])
 49
 50
 51# Load the training dataset
 52# Directory structure: root/class1/img1.jpg, root/class1/img2.jpg,
 53# root/class2/img1.jpg, ...
 54train_dataset = datasets.ImageFolder(
 55    root="dataset/cifar10/train", transform=transform)
 56
 57# Load the testing dataset that is in a similar directory structure
 58test_dataset = datasets.ImageFolder(
 59    root="dataset/cifar10/test", transform=transform)
 60
 61# Create data loaders
 62train_loader = torch.utils.data.DataLoader(
 63    dataset=train_dataset, batch_size=32, shuffle=True)
 64test_loader = torch.utils.data.DataLoader(
 65    dataset=test_dataset, batch_size=32, shuffle=False)
 66
 67# Define the network, loss function, and optimizer
 68my_model = NeuralNetwork()
 69criterion = nn.CrossEntropyLoss()
 70optimizer = torch.optim.Adam(my_model.parameters(), lr=0.001)
 71
 72# Train the neural network
 73print('##### Training the model')
 74for epoch in range(10):
 75    # Initialize the running loss for the current epoch
 76    running_loss = 0.0
 77    total_loss = 0.0
 78    # Iterate over mini-batches of training data
 79    for i, data in enumerate(train_loader, 0):
 80        inputs, labels = data
 81        # Zero the gradients to prepare for backward pass
 82        optimizer.zero_grad()
 83        outputs = my_model(inputs)
 84        # Compute the loss
 85        loss = criterion(outputs, labels)
 86        loss.backward()
 87        # Update model parameters based on computed gradients
 88        optimizer.step()
 89        running_loss += loss.item()
 90        total_loss += loss.item()
 91        if i % 200 == 199:    # Print every 200 mini-batches
 92            print(
 93                f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}"
 94            )
 95            running_loss = 0.0
 96    print(
 97        f"[{epoch + 1}] overall loss for epoch: "
 98        f"{total_loss / len(train_loader):.3f}"
 99    )
100print('Training finished')
101
102# Evaluate the neural network
103print('##### Evaluating the model')
104# Disable gradient calculation during inference
105with torch.no_grad():
106    # Initialize lists to store predicted and true labels
107    predicted_labels = []
108    true_labels = []
109    test_loss = 0.0
110    for data in test_loader:
111        # Extract inputs and labels from the data batch
112        inputs, labels = data
113        true_labels.extend(labels)
114        # Forward pass
115        outputs = my_model(inputs)
116        _, predicted = torch.max(outputs.data, 1)
117        predicted_labels.extend(predicted)
118        test_loss += criterion(outputs, labels).item()
119
120average_loss = test_loss / len(test_loader)
121print(f"Test Loss: {average_loss:.3f}")
122
123# Calculate the metrics
124metrics = ['f1-score']
125report = classification_report(true_labels, predicted_labels, output_dict=True)
126for metric in metrics:
127    metric_list = []
128    for class_label in report.keys():
129        if class_label not in ('macro avg', 'weighted avg', 'accuracy'):
130            print(f"{metric.capitalize()} for class {class_label}:",
131                  report[class_label][metric])
132            metric_list.append(report[class_label][metric])
133    metric_value = sum(metric_list) / len(metric_list)
134    print(f"Average {metric.capitalize()}: {metric_value:.2f}")
135    print(f"Accuracy: {report['accuracy']}")
136
137
138# Save the neural network
139print('##### Saving the model')
140torch.save(my_model, f"my_model_{datetime.now}.pth")
141print("The model is saved successfully")