import torch
import numpy as np
import json
import os
# Assuming `tfjs_spec` is your final dictionary containing the TensorFlow.js model specification
# Let's save it to a file named 'model.json'
def save_tfjs_model_spec(tfjs_spec, file_name="model.json"):
with open(file_name, "w") as json_file:
# Set `indent` to 4 for pretty-printing which makes the JSON file human-readable
# Ensure `ensure_ascii` is False so that the file is saved in UTF-8 encoding
json.dump(tfjs_spec, json_file, indent=4, ensure_ascii=False)
def add_conv3d_layer(
layers, filters, kernel_size, dilation, stride, layer_name
):
layers.append(
{
"class_name": "Conv3D",
"config": {
"name": layer_name,
"trainable": False,
"filters": filters,
"kernel_size": kernel_size,
"strides": stride,
"dilation_rate": dilation,
"padding": "same",
"data_format": "channels_last",
"activation": "linear", # We will add the activation layer separately
"use_bias": True,
"dtype": "float32",
},
"name": layer_name,
"inbound_nodes": [[[layers[-1]["name"], 0, 0, {}]]],
}
)
def add_activation_layer(layers, activation, layer_name):
layers.append(
{
"class_name": "Activation",
"config": {
"name": layer_name,
"trainable": False,
"dtype": "float32",
"activation": activation,
},
"name": layer_name,
"inbound_nodes": [[[layers[-1]["name"], 0, 0, {}]]],
}
)
def get_activation_name(pytorch_activation):
tfjs_activations_map = {
"ReLU": "relu",
"ELU": "elu",
"Sigmoid": "sigmoid",
"Tanh": "tanh",
"LeakyReLU": "leaky_relu",
# Add more mappings as needed
}
class_name = pytorch_activation.__class__.__name__
if class_name in tfjs_activations_map:
return tfjs_activations_map[class_name]
else:
raise ValueError(f"Unsupported PyTorch activation type: {class_name}")
def convert_pytorch_to_tfjs(pytorch_model, cube_size):
layers = [
{
"class_name": "InputLayer",
"config": {
"batch_input_shape": [None, cube_size, cube_size, cube_size, 1],
"dtype": "float32",
"sparse": False,
"ragged": False,
"name": "input",
},
"name": "input",
"inbound_nodes": [],
}
]
layer_count = 0 # Counter for naming layers
# Process each PyTorch layer and add corresponding layers to the config
for i, module in enumerate(pytorch_model.model):
if isinstance(module, torch.nn.Sequential):
for layer in module:
if isinstance(layer, torch.nn.Conv3d):
layer_name = f"conv3d_{layer_count}"
add_conv3d_layer(
layers,
layer.out_channels,
list(layer.kernel_size),
list(layer.dilation),
list(layer.stride),
layer_name,
)
else:
activation_name = f"activation_{layer_count}"
add_activation_layer(
layers, get_activation_name(layer), activation_name
)
layer_count += 1
elif isinstance(module, torch.nn.Conv3d):
layer_name = "output"
add_conv3d_layer(
layers,
module.out_channels,
list(module.kernel_size),
list(module.dilation),
list(module.stride),
layer_name,
)
layer_count += 1
# Name the last layer "output"
layers[-1]["name"] = "output"
# Generate the output model specification for TensorFlow.js
tfjs_spec = {
"format": "layers-model",
"generatedBy": "keras v2.7.0",
"convertedBy": "TensorFlow.js Converter v3.9.0",
"modelTopology": {
"keras_version": "2.6.0",
"backend": "tensorflow",
"model_config": {
"class_name": "Functional",
"config": {
"name": "model",
"layers": layers,
"input_layers": [["input", 0, 0]],
"output_layers": [[layers[-1]["name"], 0, 0]],
},
},
},
}
# Initialize the weights manifest list
weights_manifest = []
input_channels = 1
# Iterate through each layer in the layers list
for layer in layers:
class_name = layer["class_name"]
if class_name == "Conv3D":
# For Conv3D layers, we have 'kernel' (weights) and 'bias'
kernel_name = f"{layer['name']}/kernel"
bias_name = f"{layer['name']}/bias"
# Get the shape of the kernel and bias
config = layer["config"]
kernel_shape = [
config["kernel_size"][0],
config["kernel_size"][1],
config["kernel_size"][2],
input_channels,
config["filters"],
]
bias_shape = [config["filters"]]
cube_size = config["filters"] # Update cube size for the next layer
# Append the weights and bias configurations to the weights manifest
weights_manifest.extend(
[
{
"name": kernel_name,
"shape": kernel_shape,
"dtype": "float32",
},
{
"name": bias_name,
"shape": bias_shape,
"dtype": "float32",
},
]
)
# set input channels for the next layer
input_channels = config["filters"]
# Add the weightsManifest to the tfjs_spec
tfjs_spec["weightsManifest"] = [
{"paths": ["model.bin"], "weights": weights_manifest}
]
return tfjs_spec
def extract_weights(layer):
# Extract the weights and biases, move to CPU and convert to numpy arrays
layer_weight = layer.weight.data.cpu().numpy()
layer_bias = layer.bias.data.cpu().numpy()
# Transpose the weights to match TensorFlow.js's expected order
# From: (out_channels, in_channels, depth, height, width)
# To: (height, width, depth, in_channels, out_channels)
layer_weight_tfjs = np.transpose(layer_weight, (2, 3, 4, 1, 0)).flatten()
# Flatten the biases
layer_bias_tfjs = layer_bias.flatten()
return layer_weight_tfjs, layer_bias_tfjs
def save_pytorch_weights_to_bin(pytorch_model, bin_file_path="model.bin"):
# This will hold all weights and biases in the correct order
all_weights = []
# Loop through each model layer
for layer in pytorch_model.model:
# If the layer is a Sequential, recursively process its layers
if isinstance(layer, torch.nn.Sequential):
for sublayer in layer:
if isinstance(sublayer, torch.nn.Conv3d):
weights, biases = extract_weights(sublayer)
all_weights.extend([weights, biases])
elif isinstance(layer, torch.nn.Conv3d):
weights, biases = extract_weights(layer)
all_weights.extend([weights, biases])
# Concatenate all the flattened weights and biases
combined_weights = np.concatenate(all_weights)
# Convert the concatenated array to bytes
weights_bytes = combined_weights.tobytes()
# Write the bytes to a .bin file
with open(bin_file_path, "wb") as bin_file:
bin_file.write(weights_bytes)
def meshnet2tfjs(model, dirname):
# Ensure the directory exists
if not os.path.exists(dirname):
os.makedirs(dirname)
# Convert the PyTorch model to TensorFlow.js model specification
tfjs_model_spec = convert_pytorch_to_tfjs(model, 256)
# Save the TensorFlow.js model specification JSON file
model_json_path = os.path.join(dirname, "model.json")
save_tfjs_model_spec(tfjs_model_spec, model_json_path)
# Save the PyTorch model weights to a binary file
model_bin_path = os.path.join(dirname, "model.bin")
save_pytorch_weights_to_bin(model, model_bin_path)
# meshnet2tfjs(model, "/tmp/testmodel")
