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()

Parameters

  • model: The neural network model.

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

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

  • channel_last: When True, PyTorch convolutional layers permute their input and output to match the TensorFlow channel-last convention. Default False.

The filename embeds the generation type, so a PytorchGenerator invoked with generation_type="subclassing" produces pytorch_nn_subclassing.py, and generation_type="sequential" produces pytorch_nn_sequential.py.

Web Modeling Editor Support

Neural networks can also be designed visually in the BESSER Web Modeling Editor using the NNDiagram type. The backend converts the diagram into an NN metamodel instance and passes it to the PyTorch generator. From the editor’s Generate menu you can choose between the Subclassing and Sequential variants; the diagram is validated through NN.validate() before code is produced.

Output

The generated file pytorch_nn_<generation_type>.py contains:

  • Imports: PyTorch and supporting modules required by the generated code.

  • Network architecture: in subclassing mode, a NeuralNetwork(nn.Module) class with an __init__ that instantiates the layers and a forward method that chains them. In sequential mode, an nn.Sequential definition.

  • Training and evaluation block (emitted only when a Training Dataset is attached to the NN): dataset loading, loss function and optimizer setup, the training loop, evaluation against the test dataset, and saving the trained model.

The generated output for the tutorial example is shown below.

  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")