erik Hi Erik, this is code that i am using for the model. It converts image to grayscale and normalizes it
from typing import Tuple
import torch
import torch.nn as nn
class Grayscale(nn.Module):
def __init__(self, shape: Tuple[int, int, int, int], dtype=torch.float):
super(Grayscale, self).__init__()
self.shape = shape
self.dtype = dtype
def forward(self, x):
y_r = x[ :, :, 0]
y_g = x[ :, :, 1]
y_b = x[ :, :, 2]
g_ = 0.3 * y_r + 0.59 * y_g + 0.11 * y_b
g_=torch.sub(torch.div(g_,127.5),1)
return g_
def export_onnx():
shape = (32, 100, 3)
model = Grayscale(shape=shape, dtype=torch.float)
X = torch.ones(shape, dtype=torch.float)
torch.onnx.export(
model,
X,
'model2.onnx',
opset_version=12,
do_constant_folding=True
)
export_onnx()
print('done')`
