# Path: common/networks.py
# class AllConvModel(tf.Module):
#     """All convolutional network architecture."""
# 
#     def __init__(self, num_classes, num_filters, input_shape, activation=tf.nn.leaky_relu):
#         super().__init__()
#         conv_args = dict(
#             activation=activation,
#             kernel_size=3,
#             padding='same')
#         self.layers = []
#         log_resolution = int(round(
#             math.log(input_shape[0]) / math.log(2)))
#         for scale in range(log_resolution - 2):
#             self.layers.append(tf.keras.layers.Conv2D(num_filters << scale, **conv_args))
#             self.layers.append(tf.keras.layers.Conv2D(num_filters << (scale + 1), **conv_args))
#             self.layers.append(tf.keras.layers.AveragePooling2D((2, 2)))
#         self.layers.append(tf.keras.layers.Conv2D(num_classes, kernel_size=3, padding='same'))
#         self.layers.append(tf.keras.layers.GlobalAveragePooling2D())
#         self.layers.append(tf.keras.layers.Softmax())
# 
#         # call inference to instantiate variables
#         self(tf.zeros((1,) + tuple(input_shape)), training=False)
# 
#     def __call__(self, x, training=False):
#         del training  # ignore training argument since don't have batch norm
#         for layer in self.layers:
#             x = layer(x)
#         return x
# 
# class AllConvModelTorch:
#     pass
# 
# Path: common/framework.py
# def get_checkpoint_abs_path(rel_path):
#     """Converts relative path of checkpoint-related resource to absolute path.
# 
#     Args:
#         rel_path: path of the file relative to repository root.
# 
#     Returns:
#         Absolute path of the file.
# 
#     Let's say that tutorials are located in `/path/to/tutorials/` directory,
#     which means that full path of this file is `/path/to/tutorials/common/defense.py`.
#     Then following call to this method:
#         `get_checkpoint_abs_path('checkpoints/model-1')`
#     will return `/path/to/tutorials/checkpoints/model-1`
#     """
#     module_filename = inspect.getfile(inspect.currentframe())
#     module_dirname = os.path.dirname(os.path.abspath(module_filename))
#     tutorials_root = os.path.abspath(os.path.join(module_dirname, '..'))
#     return os.path.join(tutorials_root, rel_path)


import os
import sys
import tensorflow as tf
import torch
import math
import numpy as np
import logging
from common.networks import AllConvModel, AllConvModelTorch
from common.framework import get_checkpoint_abs_path
    path_torch = path_tf + ".torchmodel"
    if os.path.exists(path_torch):
        return

    print()
    print("Converting", path)
    
    # Get input sizes
    all_vars = tf.train.list_variables(
        get_checkpoint_abs_path(path_tf))

    # Is it a list of models? Or just one?
    if 'model/0' in "".join([x[0] for x in all_vars]):
        prefix = 'model/0'
    else:
        prefix = 'model'
        
    input_size, filter_size = [shape for name,shape in all_vars if prefix+'/layers/0/kernel' in name][0][2:]
    output_size = [shape for name,shape in all_vars if prefix+'/layers/9/kernel' in name][0][-1]

    num_models = sum('/0/kernel' in x for x,_ in all_vars)

    # Create the TF convnet
    convnet = [AllConvModel(num_classes=output_size,
                            num_filters=filter_size,
                            input_shape=(32, 32, input_size))
               for _ in range(num_models)]
    
    convnet_load = convnet[0] if num_models == 1 else convnet
    tf.train.Checkpoint(model=convnet_load).restore(

# Path: common/networks.py
# class AllConvModel(tf.Module):
#     """All convolutional network architecture."""
# 
#     def __init__(self, num_classes, num_filters, input_shape, activation=tf.nn.leaky_relu):
#         super().__init__()
#         conv_args = dict(
#             activation=activation,
#             kernel_size=3,
#             padding='same')
#         self.layers = []
#         log_resolution = int(round(
#             math.log(input_shape[0]) / math.log(2)))
#         for scale in range(log_resolution - 2):
#             self.layers.append(tf.keras.layers.Conv2D(num_filters << scale, **conv_args))
#             self.layers.append(tf.keras.layers.Conv2D(num_filters << (scale + 1), **conv_args))
#             self.layers.append(tf.keras.layers.AveragePooling2D((2, 2)))
#         self.layers.append(tf.keras.layers.Conv2D(num_classes, kernel_size=3, padding='same'))
#         self.layers.append(tf.keras.layers.GlobalAveragePooling2D())
#         self.layers.append(tf.keras.layers.Softmax())
# 
#         # call inference to instantiate variables
#         self(tf.zeros((1,) + tuple(input_shape)), training=False)
# 
#     def __call__(self, x, training=False):
#         del training  # ignore training argument since don't have batch norm
#         for layer in self.layers:
#             x = layer(x)
#         return x
# 
# class AllConvModelTorch:
#     pass
# 
# Path: common/framework.py
# def get_checkpoint_abs_path(rel_path):
#     """Converts relative path of checkpoint-related resource to absolute path.
# 
#     Args:
#         rel_path: path of the file relative to repository root.
# 
#     Returns:
#         Absolute path of the file.
# 
#     Let's say that tutorials are located in `/path/to/tutorials/` directory,
#     which means that full path of this file is `/path/to/tutorials/common/defense.py`.
#     Then following call to this method:
#         `get_checkpoint_abs_path('checkpoints/model-1')`
#     will return `/path/to/tutorials/checkpoints/model-1`
#     """
#     module_filename = inspect.getfile(inspect.currentframe())
#     module_dirname = os.path.dirname(os.path.abspath(module_filename))
#     tutorials_root = os.path.abspath(os.path.join(module_dirname, '..'))
#     return os.path.join(tutorials_root, rel_path)


# Path: convert_pytorch.py
import os
import sys
import tensorflow as tf
import torch
import math
import numpy as np
import logging
from common.networks import AllConvModel, AllConvModelTorch
from common.framework import get_checkpoint_abs_path

    path_torch = path_tf + ".torchmodel"
    if os.path.exists(path_torch):
        return

    print()
    print("Converting", path)
    
    # Get input sizes
    all_vars = tf.train.list_variables(
        get_checkpoint_abs_path(path_tf))

    # Is it a list of models? Or just one?
    if 'model/0' in "".join([x[0] for x in all_vars]):
        prefix = 'model/0'
    else:
        prefix = 'model'
        
    input_size, filter_size = [shape for name,shape in all_vars if prefix+'/layers/0/kernel' in name][0][2:]
    output_size = [shape for name,shape in all_vars if prefix+'/layers/9/kernel' in name][0][-1]

    num_models = sum('/0/kernel' in x for x,_ in all_vars)

    # Create the TF convnet
    convnet = [AllConvModel(num_classes=output_size,
                            num_filters=filter_size,
                            input_shape=(32, 32, input_size))
               for _ in range(num_models)]
    
    convnet_load = convnet[0] if num_models == 1 else convnet
    tf.train.Checkpoint(model=convnet_load).restore(

        get_checkpoint_abs_path(path_tf))