Customize Models

TensorFlow is deprecated. New TensorFlow model uploads are no longer accepted — author new models in PyTorch or scikit-learn. The TensorFlow mandatory variables and TensorFlow file-format examples below are retained only as reference for existing TensorFlow experiments, which remain readable.

Use Pre-trained Weights

Upload weights along with your model in the user.upload_model() step and set weights=True, the default value is False:

user.upload_model(model_name="../model-zoo/model_zoo/<task>/<framework>/model.py", weights=True)

A weights file with the same base name as the model and suffix “_weights.pkl” must exist in the same directory. For example, if the model file is “mymodel.py”, the corresponding weights file should be “mymodel_weights.pkl”.

model/
    - model.py
    - model_weights.pkl

For a pretrained model zoo template, get its weights file with user.fetch_seed_weights("<template>") first, see Upload.

Models build locally (no external hubs)

Your model file must build its architecture from local code. tracebloc does not fetch models, weights, tokenizers, or configs from an external hub (such as HuggingFace) while training — a training pod has no such egress. In practice your model file must not:

call *.from_pretrained("<hub-id>") for a model, tokenizer, or config, or
call torch.hub.load(...), or reach the same functions indirectly via getattr(X, "from_pretrained").

These are refused at upload by the model-validation check, not merely discouraged. Declaring model_id, tokenizer_id or hf_token in your model file is refused at upload too — see the note under Additional variables. To start from a pretrained model, download its weights once, build the same architecture in your model file, and upload the weights alongside it (see Use pre-trained weights). For text and other NLP tasks, also ship a tokenizer.json — pass it with user.upload_model(..., tokenizer="tokenizer.json") or place it next to your model file.

The uploaded weights file is loaded into the architecture your model file builds, matching parameter names and shapes exactly. Build the same architecture you produced the weights from, or the load will fail. The model zoo’s object detection seeds are the exception: they carry the backbone only, so the class head starts fresh from output_classes.

Model Code Formats

Model code can be:

A single Python file containing one or more functions
A single Python file with no functions
Multiple Python files (zipped), with the main file code outside of any function
Multiple Python files (zipped), with the main file code inside a function

Whatever the format, a model file may import only the packages its task allows. See Allowed imports for the list per task.

In case of multiple files saved as .zip format, the files should be directly inside the .zip file. Subdirectory structure in .zip files is not supported.

Folder models (YOLO)

A yolo object detection model trains with its own loss.py, so it is uploaded as a .zip of model.py and loss.py. user.upload_model() takes one file, and uploading model.py on its own fails with loss.py file missing in the zip. The model zoo’s YOLO templates are folders (yolo_v1/, yolo_v5/, yolo_v8/), and each ships with a ready-to-upload .zip next to it. Upload the .zip:

user.upload_model(model_name="../model-zoo/model_zoo/object_detection/pytorch/yolo_v8.zip")

For your own YOLO model, put model.py and loss.py directly in a .zip (no folder inside it) and upload that file:

my_yolo.zip
    - model.py
    - loss.py

Mandatory Variables

PyTorch mandatory variables

The file must contain these variables:

framework : name of the framework for which this model file is created. For PyTorch its value will be pytorch.
model_type : the training contract the model follows. Valid values depend on the task:
- object_detection: torchvision_detection for torchvision-style detectors (for example Faster R-CNN, RetinaNet, FCOS, SSD and the zoo’s single-file YOLO templates yolov8_s.py, yolov9_s.py, yolov10_s.py, yolo11_s.py, yolov12_s.py and yolox_s.py), or yolo for grid models that ship a loss.py (the zoo’s yolo_v1/, yolo_v5/ and yolo_v8/ folders, uploaded as a .zip, see Folder models (YOLO)). rcnn is a legacy alias of torchvision_detection. An empty or other value is refused at upload.
- keypoint_detection: direct for coordinate outputs, heatmap for heatmap outputs, or rcnn for torchvision keypoint detectors. An empty value trains as rcnn unless the model name contains “heatmap”.
- Every other PyTorch task: leave it empty.
main_class : name of the main class of the model.
image_size : name of the variable defining the input size considered for creating model layers.
category : name of the category for which this model file is created. PyTorch supports every task: image_classification, object_detection, keypoint_detection, semantic_segmentation, text_classification, sentence_pair_classification, token_classification, masked_language_modeling, causal_language_modeling, seq2seq, embeddings, tabular_classification, tabular_regression, time_series_classification, time_series_forecasting, time_to_event_prediction. See Model parameters by task for the variables each task adds.
batch_size : name of the variable defining the batch size of the images considered for creating model layers.

Sklearn mandatory variables

The file must contain these variables:

framework : name of the framework for which this model file is created. For Sklearn its value will be sklearn.
model_type : name of the model type for which this model file is created. Its value will be either empty or tree or linear or knn or naive or mlp.
main_class or main_method : name of the main class or main method of the model. Any one of these is necessary for model.
image_size : name of the variable defining the input size considered for creating model layers.
category : name of the category for which this model file is created. Sklearn supports tabular_classification and tabular_regression. For time_to_event_prediction, use the lifelines or scikit_survival framework, or PyTorch.
batch_size : name of the variable defining the batch size of the images considered for creating model layers.

Tensorflow mandatory variables

Deprecated — retained for existing TensorFlow experiments. New uploads must use PyTorch or scikit-learn.

Each format must contain these variables on the main file:

framework : name of the framework for which this model file is created. For tensorflow its value will be tensorflow.
model_type : name of the model type for which this model file is created. Its value will be either empty or rcnn or heatmap.
main_method : name of the main function
input_shape or image_size : name of the variable defining the input shape of the model
category : name of the category for which this model file is created. Its value will be either of these : image_classification, object_detection, keypoint_detection, text_classification.
output_classes : name of the variable defining the output shape of the variables. The value of this variable is equal to the number of classes in the dataset.

Additional variables

Some additional variables are required for specific categories

num_feature_points : number of keypoints (keypoint_detection) or of input columns (tabular_classification, tabular_regression, time_to_event_prediction) for which this model file is created. This variable is used only for those categories.

Do not declare model_id, tokenizer_id or hf_token in your model file. If any of them is set, the SDK refuses the upload with “HuggingFace / hub-referenced models are no longer supported”. A training pod has no egress to an external hub, so a hub id could not be fetched anyway (see Models build locally (no external hubs)). Build the architecture locally and upload its weights instead. Hugging Face support is removed.

The framework variable is compulsory and should always be placed at the top of your code just after the imports, before any other variable.
Changing the sequence of the framework variable will cause the model upload to fail.
Do not add any comment after these variables. It will cause model upload to fail.

Detailed Information on each format

1. Single python file containing one class for pyTorch

All your model code needs to be contained in a single python file with the following structure:

1.1 Import Section

Import all the necessary functionality needed to build your model.

# Example
import torch
import torch.nn as nn

1.2 Main Class

This is the main method which should return the model as output.

#Example
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
            nn.BatchNorm2d(6),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(16 * 53 * 53, 120)
        self.relu = nn.ReLU()
        self.fc1 = nn.Linear(120, 84)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(84, 1)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        out = self.relu(out)
        out = self.fc1(out)
        out = self.relu1(out)
        out = self.fc2(out)
        return out

The value of the above variables would look like this:

framework = ‘pytorch’
main_method = ‘Net’
image_size = 224
batch_size = 16

An example model file of this format can be found here.

2. Single python file containing one class for sklearn

All your model code needs to be contained in a single python file with the following structure:

2.2 Import Section

Import all the necessary functionality you need to build your model.

# Example
from sklearn.svm import SVC

2.3 Main Method

This is the main method which should return the model as output.

#Example
def MyModel():
    return SVC(kernel='linear')

The value of the above variables would look like this:

framework = ‘sklearn’
main_method = ‘MyModel’
image_size = 224
batch_size = 16

3. Single python file containing one/multiple methods for tensorflow

Formats 3 and 4 (TensorFlow) are deprecated and retained only for existing TensorFlow experiments. New uploads must use the PyTorch or scikit-learn formats above.

All your model code needs to be contained in a single python file with the following structure:

3.1 Import Section

Import all the necessary functionality you need to build your model.

# Example
import tensorflow as tf
from tensorflow.keras import layers, utils

3.2 Supporting Functions

Any kind of functionality that supports the model architecture. This section is optional. This section can be left out if the model architecture does not require it.

def dense_block(input_tensor, k, block_reps):
	"""
	tensor: input tensor from the previous layers
	k: growth rate
	block_reps: Number of times the block is repeated
	Return the concatenated tensors
	"""
	for _ in range(block_reps):

		x = layers.BatchNormalization()(input_tensor)
		x = layers.ReLU()(x)
		x = layers.Conv2D(filters=4*k, kernel_size=1)(x)
		x = layers.BatchNormalization()(x)
		x = layers.ReLU()(x)
		x = layers.Conv2D(filters=k, kernel_size=3, padding='same')(x)
		output_tensor = layers.Concatenate()([input_tensor, x])

		return output_tensor

3.3 Main Method

This is the main method which should return the model as output.

#Example
def MyModel(input_shape=(224,224,3),classes=3):

	k = 32 #growth rate

	input = layers.Input(shape=input_shape)
	x = layers.BatchNormalization()(input)
	x = layers.ReLU()(x)
	x = layers.Conv2D(filters=2*k, kernel_size=7, strides=2)(x)
	x = layers.MaxPooling2D(pool_size=3, strides=2)(x)

	x = dense_block(x, 32, 6)
	x = transition_layers(x)

	x = dense_block(x, 32, 12)
	x = transition_layers(x)

	x = dense_block(x, 32, 32)
	x = transition_layers(x)

	x = dense_block(x, 32, 32)

	x = layers.GlobalAveragePooling2D()(x)
	output = layers.Dense(classes, activation='softmax')(x)

	model_densenet = tf.keras.Model(input, output)

	return model_densenet

The value of the above variables would look like this:

framework = ‘tensorflow’
main_method = ‘MyModel’
input_shape = ‘input_shape’
output_classes = ‘classes’

4. Single python file having no method for tensorflow

All your model code needs to be in a single python file with the following structure: The model file can be broken down into two sections: The Import section and the main code.

4.1 Import Section

Import all the necessary functionalities you need to build your model.

# Example
import tensorflow as tf
from tensorflow.keras import layers, models

4.2 Main Code

This is the model code with all layers.

#Example
input_shape=(224,224,3)
classes=3

model = models.Sequential()
# layer conv 1
model.add(layers.Conv2D(32, 3, activation='relu', input_shape=input_shape))
model.add(layers.MaxPooling2D(2))
layers.BatchNormalization()
# layer conv 2
model.add(layers.Conv2D(64, 3, activation='relu', input_shape=input_shape))
model.add(layers.MaxPooling2D(2))
layers.BatchNormalization()
# layer conv 3
model.add(layers.Conv2D(128, 3, activation='relu', input_shape=input_shape))
model.add(layers.MaxPooling2D(2))
layers.BatchNormalization()
# layer conv 4
model.add(layers.Conv2D(264, 3, activation='relu', input_shape=input_shape))
model.add(layers.MaxPooling2D(2))
layers.BatchNormalization()
# layer conv 5
model.add(layers.Conv2D(1, 1, activation='relu', input_shape=input_shape))
layers.BatchNormalization()
# Flatten the feature maps to serve dense
model.add(layers.Flatten())
model.add(layers.Dense(classes, activation='softmax'))

The value of the above variables would look like this:

framework = ‘tensorflow’
main_method = ”
input_shape = ‘input_shape’
output_classes = ‘classes’

Add the mandatory variables at the top of the file (after imports) to get a complete model file ready for upload.

5. Multiple python files with main file code not contained in a method

In this format the model architecture is defined using multiple python files. But the main file should use the same format as type 1. In order to upload models with this format, zip all python files with the main method containing the mentioned variables and upload the zip file.

The main file should contain the variable values as mentioned in the type 1 format. Other files should not contain these variables.

6. Multiple python file with main file code contained in a method

In this format the model architecture is defined using multiple python files. But the main file should use the same format as type 2. In order to upload models with this format, zip all python files with the main method containing the mentioned variables and upload the zip file.

The main file should contain the variable values as mentioned in type 2 format. Other files should not contain these variables.

Next Steps

Configure training parameters: Hyperparameters

Need Help?

Email us at [email protected]

Overview

Create a Use Case

Join a Use Case

Environment Setup

Tools & Help

Releases

Use Pre-trained Weights

Models build locally (no external hubs)

Model Code Formats

Folder models (YOLO)

Mandatory Variables

PyTorch mandatory variables

Sklearn mandatory variables

Tensorflow mandatory variables

Additional variables

Detailed Information on each format

1. Single python file containing one class for pyTorch

1.1 Import Section

1.2 Main Class

2. Single python file containing one class for sklearn

2.2 Import Section

2.3 Main Method

3. Single python file containing one/multiple methods for tensorflow

3.1 Import Section

3.2 Supporting Functions

3.3 Main Method

4. Single python file having no method for tensorflow

4.1 Import Section

4.2 Main Code

5. Multiple python files with main file code not contained in a method

6. Multiple python file with main file code contained in a method

Next Steps

Need Help?

​Use Pre-trained Weights

​Models build locally (no external hubs)

​Model Code Formats

​Folder models (YOLO)

​Mandatory Variables

​PyTorch mandatory variables

​Sklearn mandatory variables

​Tensorflow mandatory variables

​Additional variables

​Detailed Information on each format

​1. Single python file containing one class for pyTorch

​1.1 Import Section

​1.2 Main Class

​2. Single python file containing one class for sklearn

​2.2 Import Section

​2.3 Main Method

​3. Single python file containing one/multiple methods for tensorflow

​3.1 Import Section

​3.2 Supporting Functions

​3.3 Main Method

​4. Single python file having no method for tensorflow

​4.1 Import Section

​4.2 Main Code

​5. Multiple python files with main file code not contained in a method

​6. Multiple python file with main file code contained in a method

​Next Steps

​Need Help?

Use Pre-trained Weights

Models build locally (no external hubs)

Model Code Formats

Folder models (YOLO)

Mandatory Variables

PyTorch mandatory variables

Sklearn mandatory variables

Tensorflow mandatory variables

Additional variables

Detailed Information on each format

1. Single python file containing one class for pyTorch

1.1 Import Section

1.2 Main Class

2. Single python file containing one class for sklearn

2.2 Import Section

2.3 Main Method

3. Single python file containing one/multiple methods for tensorflow

3.1 Import Section

3.2 Supporting Functions

3.3 Main Method

4. Single python file having no method for tensorflow

4.1 Import Section

4.2 Main Code

5. Multiple python files with main file code not contained in a method

6. Multiple python file with main file code contained in a method

Next Steps

Need Help?