import torch

import torch.nn as nn

import traceback

from Caffe import caffe_net

import torch.nn.functional as F

from torch.autograd import Variable

from Caffe import layer_param

from torch.nn.modules.utils import _pair

import numpy as np

import inspect

"""

How to support a new layer type:

layer_name=log.add_layer(layer_type_name)

top_blobs=log.add_blobs(<output of that layer>)

layer=caffe_net.Layer_param(xxx)

<set layer parameters>

[<layer.add_data(*datas)>]

log.cnet.add_layer(layer)

Please MUTE the inplace operations to avoid not find in graph

"""

NET_INITTED=False

WARNING_STRINGS=''

RP_TRANSFERRING_FLAG=False # this flag prevents transferring Rp function in Rp function.

class TransLog(object):

self._blobs[id(new_tensor)]=blob_name

log=TransLog()

layer_names={}

class Rp(object):

return out

# ----- for torch.nn.functional operations -----

def _conv2d(raw,input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):

return x

def _conv_transpose2d(raw,input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):

return x

def _interpolate(raw,input, size=None, scale_factor=None, mode='nearest', align_corners=None):

return x

def _linear(raw,input, weight, bias=None):

return x

def _pool(type,raw,input,x,kernel_size,stride,padding,ceil_mode):

log.cnet.add_layer(slice2_layer)

def _max_pool2d(raw,input, kernel_size, stride=None, padding=0, dilation=1,

return x

def _avg_pool2d(raw,input, kernel_size, stride = None, padding = 0, ceil_mode = False, count_include_pad = True):

x = raw(input, kernel_size, stride, padding, ceil_mode, count_include_pad)

_pool('ave',raw,input, x, kernel_size, stride, padding,ceil_mode)

return x

def _dropout(raw,input,p=0.5, training=False, inplace=False):

return x

def _threshold(raw,input, threshold, value, inplace=False):

return x

def _relu(raw, input, inplace=False):

return x

def _prelu(raw, input, weight):

return x

def _leaky_relu(raw, input, negative_slope=0.01, inplace=False):

return x

def _tanh(raw, input):

return x

def _softmax(raw, input, dim=None, _stacklevel=3):

return x

def _batch_norm(raw,input, running_mean, running_var, weight=None, bias=None,

return x

def _instance_norm(raw, input, running_mean=None, running_var=None, weight=None,

return x

def op_placeholder(raw, *args, **kwargs):

return output

F_supported=[

]

for op_name in F.__dict__:

if op_name in F_supported:

raw_func=getattr(F, op_name)

transfer_func=globals()['_'+op_name]

op_wrapper=Rp(raw_func,transfer_func)

setattr(F, op_name, op_wrapper)

else:

if op_name[0]=='_' or op_name in ['division','warnings','math','torch','utils','vision','Col2Im','Im2Col','grad','weak_script','List']:

continue

setattr(F,op_name,Rp(getattr(F,op_name),op_placeholder))

# ----- for torch operations -----

def torch_max(raw,*args):

return x

def torch_cat(raw,inputs, dimension=0):

return x

def torch_split(raw,tensor, split_size, dim=0):

return x

def torch_add(raw,*args):

return x

def torch_sub(raw,*args):

return ___sub__(*args)

def torch_mul(raw,*args):

return ___mul__(*args)

def torch_div(raw,*args):

return ___div__(*args)

def torch_pow(raw,*args):

return x

def torch_sqrt(raw,*args):

return x

torch_op_supported=[

]

for op_name in torch_op_supported:

raw_op = getattr(torch, op_name)

op_wrapper=Rp(raw_op,globals()['torch_'+op_name])

setattr(torch, op_name, op_wrapper)

# ----- for Variable/torch.Tensor operations --------

def _view(input, *args):

return x

def _mean(input, *args,**kwargs):

return x

def _sum(input, *args,**kwargs):

return x

def _contiguous(input,*args):

x=raw_tensor_magic_op['contiguous'](input,*args)

log.reuse_blob(input,x)

return x

def _add(input,*args):

return ___add__(input, *args)

def _sub(input,*args):

return ___sub__(input, *args)

def _mul(input,*args):

return ___mul__(input, *args)

def _div(input,*args):

return ___div__(input, *args)

def _pow(input,*args):

return ___pow__(input, *args)

def _sqrt(input, *args):

return x

def ___add__(input, *args):

return x

def ___iadd__(input, *args):

return x

def ___sub__(input, *args):

return x

def ___isub__(input, *args):

return x

def ___mul__(input, *args):

return x

def ___imul__(input, *args):

return x

def ___div__(input, *args):

return x

def ___pow__(input, *args):

return x

# TODO: other types of the view function

tensor_op_supported=[]

tensor_magic_op_supported=[

]

raw_tensor_magic_op={}

if hasattr(Variable,'__add__'):

tensor_target=Variable

else:

# for new version >=0.4.0

tensor_target=torch.Tensor

for op_name in tensor_magic_op_supported:

raw_op=getattr(tensor_target,op_name)

raw_tensor_magic_op[op_name]=raw_op

setattr(tensor_target,op_name,globals()['_'+op_name])

for op_name in tensor_op_supported:

raw_op = getattr(tensor_target, op_name)

op_wrapper = Rp(raw_op, globals()['_' + op_name])

setattr(tensor_target, op_name, op_wrapper)

def trans_net(net,input_var,name='TransferedPytorchModel'):

print(WARNING_STRINGS)

def save_prototxt(save_name):

log.cnet.save_prototxt(save_name)

def save_caffemodel(save_name):

log.cnet.save(save_name)