zhutian
30 lines
782 B
Python
30 lines
782 B
Python
import copy
|
|
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
from onnx_tf.handlers.backend_handler import BackendHandler
|
|
from onnx_tf.handlers.handler import onnx_op
|
|
from onnx_tf.handlers.handler import tf_func
|
|
|
|
|
|
@onnx_op("LRN")
|
|
@tf_func(tf.nn.lrn)
|
|
class LRN(BackendHandler):
|
|
|
|
@classmethod
|
|
def version_1(cls, node, **kwargs):
|
|
attrs = copy.deepcopy(node.attrs)
|
|
alpha = attrs.get("alpha", 1e-4)
|
|
attrs.setdefault("beta", 0.75)
|
|
size = attrs["size"]
|
|
attrs["alpha"] = alpha / size
|
|
attrs["depth_radius"] = np.floor([(size - 1) / 2.])[0]
|
|
# TODO: LRN in tf accepts radius
|
|
# but in ONNX/Caffe accepts diameter.
|
|
# This could be a problem.
|
|
return [
|
|
cls.make_tensor_from_onnx_node(
|
|
node, attrs=attrs, c_last_only=True, **kwargs)
|
|
]
|