mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 09:36:08 +00:00
Simple onnx.py clean ups (#8054)
* start * simplify ops * why did this not work before * will split buffer parse to separate pr * flip the error order * only this much for now * to_python_const clean up * minimize diff * move tensor_methods into onnx.py * improve some type signatures --------- Co-authored-by: chenyu <[email protected]>
This commit is contained in:
geohotstan
+77
-118
@@ -1,12 +1,10 @@
|
||||
from __future__ import annotations
|
||||
from typing import List, Dict, Union
|
||||
import importlib
|
||||
from functools import lru_cache
|
||||
from typing import List, Dict, Union, Callable, Any
|
||||
import importlib, functools
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, dtypes, Device
|
||||
from tinygrad.tensor import _to_np_dtype
|
||||
from tinygrad.helpers import getenv, DEBUG, CI, OSX
|
||||
from tinygrad.dtype import ConstType, DType
|
||||
from tinygrad import Tensor, dtypes
|
||||
from tinygrad.helpers import getenv, DEBUG
|
||||
from tinygrad.dtype import DType, ConstType
|
||||
from tinygrad.device import is_dtype_supported
|
||||
from onnx import AttributeProto, ModelProto, TensorProto, TypeProto
|
||||
try:
|
||||
@@ -17,32 +15,56 @@ except ImportError:
|
||||
def tensor_dtype_to_np_dtype(tensor_dtype:int) -> np.dtype: return TENSOR_TYPE_TO_NP_TYPE[tensor_dtype]
|
||||
|
||||
cache_misses = 0
|
||||
@lru_cache(None)
|
||||
def _cached_to_python_const(t:Tensor, tobytes): return t.data().tobytes() if tobytes else t.tolist()
|
||||
@functools.lru_cache(None)
|
||||
def _cached_to_python_const(t:Tensor):
|
||||
if t.dtype is dtypes.uint8: return t.data().tobytes()
|
||||
if 0 in t.shape: return []
|
||||
return t.tolist()
|
||||
|
||||
# Tensor -> python value cache for parameters
|
||||
def to_python_const(t, tobytes=False) -> Union[List[ConstType], List[bytes], Union[ConstType, bytes]]:
|
||||
def to_python_const(t) -> Union[List[ConstType], List[bytes], Union[ConstType, bytes]]:
|
||||
if not isinstance(t, Tensor): return t
|
||||
global cache_misses
|
||||
ret = _cached_to_python_const(t, tobytes)
|
||||
ret = _cached_to_python_const(t)
|
||||
if (info := _cached_to_python_const.cache_info()).misses > cache_misses and DEBUG >= 3:
|
||||
print(f"Cache miss for {t}, {tobytes=}")
|
||||
print(f"Cache miss for {t}")
|
||||
cache_misses = info.misses
|
||||
return ret
|
||||
|
||||
# src: onnx/mapping.py https://onnx.ai/onnx/api/mapping.html#l-mod-onnx-mapping
|
||||
# not supported: STRING = 8 COMPLEX64 = 14, COMPLEX128 = 15, UINT4 = 21, INT4 = 22
|
||||
# TODO: use dtypes.float16 for FLOAT16
|
||||
DTYPE_MAP: Dict[TensorProto.DataType, DType] = {
|
||||
TensorProto.FLOAT:dtypes.float, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8, TensorProto.UINT16:dtypes.uint16,
|
||||
TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64, TensorProto.BOOL:dtypes.bool,
|
||||
TensorProto.FLOAT16:dtypes.float, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32, TensorProto.UINT64:dtypes.uint64,
|
||||
TensorProto.BFLOAT16:dtypes.bfloat16, TensorProto.FLOAT8E4M3FN:dtypes.float, TensorProto.FLOAT8E4M3FNUZ:dtypes.float,
|
||||
TensorProto.FLOAT8E5M2:dtypes.float, TensorProto.FLOAT8E5M2FNUZ:dtypes.float
|
||||
# TODO: use real float16
|
||||
# src: onnx/mapping.py
|
||||
DTYPE_MAP: Dict[TensorProto.DataType | int, DType] = {
|
||||
TensorProto.FLOAT:dtypes.float32, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8,
|
||||
TensorProto.UINT16:dtypes.uint16, TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64,
|
||||
TensorProto.BOOL:dtypes.bool, TensorProto.FLOAT16:dtypes.float32, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32,
|
||||
TensorProto.UINT64:dtypes.uint64, TensorProto.BFLOAT16:dtypes.bfloat16, TensorProto.FLOAT8E4M3FN:dtypes.float,
|
||||
TensorProto.FLOAT8E4M3FNUZ:dtypes.float, TensorProto.FLOAT8E5M2:dtypes.float, TensorProto.FLOAT8E5M2FNUZ:dtypes.float
|
||||
}
|
||||
def dtype_parse(onnx_dtype: TensorProto.DataType | int) -> DType:
|
||||
if onnx_dtype not in DTYPE_MAP: raise NotImplementedError(f"onnx dtype {TensorProto.DataType.Name(onnx_dtype)} is not supported")
|
||||
return DTYPE_MAP[onnx_dtype] if is_dtype_supported(DTYPE_MAP[onnx_dtype]) else dtypes.float
|
||||
|
||||
# src: onnx/onnx_ml_pb2.pyi
|
||||
ATTRIBUTE_MAP: Dict[AttributeProto.AttributeType, Callable[[AttributeProto], Any]] = {
|
||||
AttributeProto.FLOAT: lambda a: float(a.f), AttributeProto.INT: lambda a: int(a.i),
|
||||
AttributeProto.STRING: lambda a: a.s.decode("utf-8"), AttributeProto.TENSOR: lambda a: buffer_parse(a.t),
|
||||
AttributeProto.FLOATS: lambda a: tuple(float(x) for x in a.floats), AttributeProto.INTS: lambda a: tuple(int(x) for x in a.ints),
|
||||
AttributeProto.STRINGS: lambda a: tuple(x.decode("utf-8") for x in a.strings)
|
||||
}
|
||||
def attribute_parse(onnx_attribute: AttributeProto):
|
||||
if onnx_attribute.type not in ATTRIBUTE_MAP:
|
||||
raise NotImplementedError(f"attribute with type {AttributeProto.AttributeType.Name(onnx_attribute.type)} is not supported")
|
||||
return ATTRIBUTE_MAP[onnx_attribute.type](onnx_attribute)
|
||||
|
||||
def buffer_parse(inp: TensorProto) -> Tensor:
|
||||
if dat := list(inp.float_data) or list(inp.int32_data) or list(inp.int64_data):
|
||||
return Tensor(dat, dtype=dtype_parse(inp.data_type), requires_grad=False).reshape(tuple(inp.dims))
|
||||
if len(inp.raw_data) > 0:
|
||||
return Tensor(np.frombuffer(inp.raw_data, dtype=tensor_dtype_to_np_dtype(inp.data_type)).copy().reshape(tuple(inp.dims)),
|
||||
dtype=dtype_parse(inp.data_type), requires_grad=False)
|
||||
raise NotImplementedError(f"buffer with data type {TensorProto.DataType.Name(inp.data_type)} is not supported")
|
||||
|
||||
onnx_ops = importlib.import_module('extra.onnx_ops')
|
||||
|
||||
ONNXLIMIT = getenv("ONNXLIMIT", -1)
|
||||
|
||||
def get_run_onnx(onnx_model: ModelProto):
|
||||
@@ -66,54 +88,30 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
elif attr == 'sparse_tensor_type': raise NotImplementedError(f"sparse_tensor_type is not implemented: {type_proto}")
|
||||
else: raise AttributeError(f"unknown attr: {attr}, {type_proto}")
|
||||
|
||||
def buffer_parse(inp: TensorProto) -> Tensor:
|
||||
if inp.data_type not in DTYPE_MAP:
|
||||
raise NotImplementedError(f"data type not supported {inp.name} {inp.dims} {inp.data_type}")
|
||||
dtype = DTYPE_MAP[inp.data_type] if is_dtype_supported(DTYPE_MAP[inp.data_type]) else dtypes.float32
|
||||
if dat := list(inp.float_data) or list(inp.int32_data) or list(inp.int64_data):
|
||||
return Tensor(dat, dtype=dtype, requires_grad=False).reshape(tuple(inp.dims))
|
||||
if len(inp.raw_data) > 0:
|
||||
data = np.frombuffer(inp.raw_data, dtype=tensor_dtype_to_np_dtype(inp.data_type)).astype(_to_np_dtype(dtype)).copy()
|
||||
return Tensor(data.reshape(tuple(inp.dims)), requires_grad=False)
|
||||
return Tensor(None, requires_grad=False)
|
||||
|
||||
def attribute_parse(a: AttributeProto) -> float | int | str | Tensor | tuple[float] | tuple[int]:
|
||||
# TODO: this is not complete, see onnx/onnx_ml_pb2.pyi for a complete list
|
||||
if a.type == AttributeProto.FLOAT: return float(a.f)
|
||||
elif a.type == AttributeProto.INT: return int(a.i)
|
||||
elif a.type == AttributeProto.STRING: return a.s.decode("utf-8")
|
||||
elif a.type == AttributeProto.TENSOR: return buffer_parse(a.t) # TENSOR
|
||||
elif a.type == AttributeProto.FLOATS: return tuple(float(x) for x in a.floats)
|
||||
elif a.type == AttributeProto.INTS: return tuple(int(x) for x in a.ints)
|
||||
elif a.type == AttributeProto.STRINGS: return tuple(x.decode("utf-8") for x in a.strings)
|
||||
elif a.type == AttributeProto.GRAPH: raise NotImplementedError(f"graph not implemented: {a.g}\n likely an OP requiring control flow")
|
||||
else: raise RuntimeError(f"can't parse {a.type} {a}")
|
||||
|
||||
tensors: Dict[str, Tensor] = {}
|
||||
|
||||
# get weights and biases
|
||||
for inp in onnx_model.graph.initializer:
|
||||
tensors[inp.name] = buffer_parse(inp)
|
||||
|
||||
# preparse the attributes
|
||||
attribute_dict = {}
|
||||
domain = ""
|
||||
for num,n in enumerate(onnx_model.graph.node):
|
||||
attribute_dict[num] = {x.name:attribute_parse(x) for x in n.attribute}
|
||||
if n.domain: domain = n.domain
|
||||
# initialization data
|
||||
model_parameters = {inp.name:buffer_parse(inp) for inp in onnx_model.graph.initializer}
|
||||
model_attributes = {num:{x.name:attribute_parse(x) for x in n.attribute} for num,n in enumerate(onnx_model.graph.node)}
|
||||
|
||||
# model specs
|
||||
is_onnx_preview_training = any(n.HasField("domain") and n.domain == "ai.onnx.preview.training" for n in onnx_model.graph.node)
|
||||
onnx_model_version = onnx_model.opset_import[0].version
|
||||
|
||||
# mapping from onnx ops to tensor.py ops
|
||||
tensor_methods = {
|
||||
op:op.lower() for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan", "Asin", "Acos", "Atan",
|
||||
"Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsInf", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh", "Tanh",
|
||||
"Softsign", "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Xor", "Round", "Erf")
|
||||
}
|
||||
|
||||
def run_onnx(inputs={}, debug=0):
|
||||
debug = getenv("DEBUGONNX") or debug
|
||||
input_tensors: Dict[str,Tensor|List[Tensor]] = {}
|
||||
intermediate_tensors: Dict[str,Tensor] = {}
|
||||
output_tensor_names = [x.name for x in onnx_model.graph.output]
|
||||
|
||||
# get inputs
|
||||
for model_input in onnx_model.graph.input:
|
||||
name = model_input.name
|
||||
if name in tensors: continue
|
||||
if name in model_parameters: continue
|
||||
shape = type_parse(model_input.type)
|
||||
if name in inputs:
|
||||
if isinstance(inputs[name], Tensor):
|
||||
@@ -121,7 +119,7 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
elif isinstance(inputs[name], list):
|
||||
input_tensors[name] = [Tensor(i, requires_grad=False) for i in inputs[name]]
|
||||
# TODO: this is just to make training tests pass, need a principled way to handle training vs non-training
|
||||
elif domain == "ai.onnx.preview.training":
|
||||
elif is_onnx_preview_training:
|
||||
input_tensors[name] = Tensor(inputs[name], requires_grad=True)
|
||||
else:
|
||||
input_tensors[name] = Tensor(inputs[name], requires_grad=False)
|
||||
@@ -133,62 +131,27 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
raise RuntimeError(f"no data for {name} with shape {shape}")
|
||||
|
||||
def fetch_tensor(x: str):
|
||||
if x in tensors: return tensors[x]
|
||||
if x in model_parameters: return model_parameters[x]
|
||||
if x in intermediate_tensors: return intermediate_tensors[x]
|
||||
if x != "": return input_tensors[x]
|
||||
return None
|
||||
|
||||
for num,n in enumerate(onnx_model.graph.node):
|
||||
inp: List[Tensor] = []
|
||||
if debug >= 3: print("inputs:")
|
||||
for x in n.input:
|
||||
t = fetch_tensor(x)
|
||||
if debug >= 3: print(f"\t{x} - {t}")
|
||||
inp.append(t)
|
||||
opt: Dict = attribute_dict[num]
|
||||
if debug >= 1: print(f"{num}: op {n.op_type} shape {[x.shape if isinstance(x, Tensor) else x for x in inp]} opt {opt}")
|
||||
inp = [fetch_tensor(x) for x in n.input]
|
||||
opt = model_attributes[num]
|
||||
|
||||
if debug >= 1: print(f"{num}: op \"{n.op_type}\" input shapes {[x.shape if isinstance(x, Tensor) else x for x in inp]} opt {opt}")
|
||||
if debug >= 3: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {t}" for i,(x,t) in enumerate(zip(n.input, inp))))
|
||||
|
||||
if n.op_type in tensor_methods:
|
||||
ret = getattr(Tensor, tensor_methods[n.op_type])(*inp, **opt)
|
||||
|
||||
# NOTE some ops live here because they require access to some local variables
|
||||
# have to use n.output for cases when num_outputs is absent
|
||||
if n.op_type in onnx_ops.tensor_methods:
|
||||
ret = getattr(Tensor, n.op_type.lower())(*inp, **opt)
|
||||
elif n.op_type == "Split":
|
||||
axis = opt.get("axis", 0)
|
||||
split = None if len(inp) == 1 else to_python_const(inp[1])
|
||||
if split is None:
|
||||
split = [inp[0].shape[axis] // len(n.output)] * len(n.output)
|
||||
for i in range(inp[0].shape[axis] % len(n.output)):
|
||||
split[i] += 1
|
||||
i, ret = 0, []
|
||||
arg = [None] * inp[0].ndim
|
||||
for s in split:
|
||||
arg[axis] = (i,i+s)
|
||||
ret.append(inp[0].shrink(arg=tuple(arg)))
|
||||
i = i+s
|
||||
ret = tuple(ret)
|
||||
|
||||
# need to check onnx_model_version
|
||||
elif n.op_type == "Slice":
|
||||
if onnx_model_version < 10:
|
||||
axes, ends, starts, steps = list(opt.get("axes", range(inp[0].ndim))), list(opt["ends"]), list(opt["starts"]), [1]*inp[0].ndim
|
||||
else:
|
||||
starts, ends = inp[1:3]
|
||||
axes = list(range(inp[0].ndim)) if len(inp) <= 3 else to_python_const(inp[3].cast(dtypes.int32))
|
||||
steps = inp[4].cast(dtypes.int32).tolist() if len(inp) > 4 else [1]*inp[0].ndim
|
||||
starts, ends = to_python_const(starts), to_python_const(ends)
|
||||
arg = [(0,x,1) for x in inp[0].shape]
|
||||
for i, axis in enumerate(axes):
|
||||
axis = int(axis) + inp[0].ndim if axis < 0 else int(axis)
|
||||
if starts[i] < 0: starts[i] += inp[0].shape[axis]
|
||||
if ends[i] < 0: ends[i] += inp[0].shape[axis]
|
||||
starts[i], ends[i] = max(0, min(starts[i], inp[0].shape[axis])), max(0, min(ends[i], inp[0].shape[axis]))
|
||||
if starts[i] > ends[i] and steps[i] >= 0: steps[i] = -steps[i]
|
||||
arg[axis] = (starts[i], ends[i], steps[i])
|
||||
new_shape = tuple((s, e) if st > 0 else (e+1, s+1) for s, e, st in arg)
|
||||
if any(s==e for s,e in new_shape): ret = inp[0].shrink(new_shape)
|
||||
else: ret = inp[0][tuple([slice(s,e,st) for s,e,st in arg])]
|
||||
|
||||
# need to call backward on intermediate_tensors
|
||||
axis, n_outputs = opt.get('axis', 0), opt.get('num_outputs') or len(n.output)
|
||||
sz = inp[0].shape[axis]
|
||||
sizes = to_python_const(inp[1]) if len(inp) == 2 else [sz // n_outputs + (1 if i < sz % n_outputs else 0) for i in range(n_outputs)]
|
||||
ret = inp[0].split(sizes, axis)
|
||||
elif n.op_type == "Gradient":
|
||||
assert len(opt["xs"]) == len(inp), f"len(opt['xs']):{len(opt['xs'])}, len(inp):{len(inp)} output and input has to match"
|
||||
y = opt["y"]
|
||||
@@ -209,16 +172,12 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
print("UNSUPPORTED", n.op_type, n.input, n.output)
|
||||
raise NotImplementedError(f"op_type {n.op_type} not supported")
|
||||
|
||||
# finalization after running the op
|
||||
if not isinstance(ret, tuple): ret = (ret, )
|
||||
assert len(n.output) <= len(ret), f"expected output size must be less than {len(ret)}, it's {n.output}"
|
||||
if debug >= 2: print([x.shape if isinstance(x, Tensor) else None for x in ret])
|
||||
if debug >= 2: print("outputs:")
|
||||
for i in range(len(n.output)):
|
||||
if debug >= 2: print(f"\t{n.output[i]} - {ret[i]}")
|
||||
intermediate_tensors[n.output[i]] = ret[i]
|
||||
if num == ONNXLIMIT:
|
||||
output_tensor_names = n.output
|
||||
break
|
||||
if len(n.output) > len(ret): raise RuntimeError(f"expected output size must be less than {len(ret)}, it's {n.output}")
|
||||
for i in range(len(n.output)): intermediate_tensors[n.output[i]] = ret[i]
|
||||
if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{n.output[i]} - {ret[i]}" for i in range(len(n.output))))
|
||||
|
||||
return {outp:intermediate_tensors[outp] for outp in output_tensor_names}
|
||||
if num == ONNXLIMIT: return {name:intermediate_tensors[name] for name in n.output}
|
||||
return {x.name:intermediate_tensors[x.name] for x in onnx_model.graph.output}
|
||||
return run_onnx
|
||||
|
||||
+12
-8
@@ -3,13 +3,9 @@ from typing import Union, Tuple, Optional, List, Any, cast
|
||||
from tinygrad.tensor import Tensor, _broadcast_shape
|
||||
from tinygrad.dtype import ImageDType, dtypes
|
||||
from tinygrad.helpers import prod, flatten
|
||||
from extra.onnx import DTYPE_MAP, to_python_const
|
||||
from extra.onnx import dtype_parse, to_python_const
|
||||
import numpy as np
|
||||
|
||||
tensor_methods = {"Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan", "Asin", "Acos", "Atan","Relu",
|
||||
"Sigmoid", "MatMul", "Floor", "Ceil", "IsInf", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh", "Tanh", "Softsign",
|
||||
"Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Xor", "Round", "Erf"}
|
||||
|
||||
# **************** Free Ops ****************
|
||||
|
||||
def Identity(x: Tensor): return x
|
||||
@@ -26,7 +22,7 @@ def Min(*data_0): return functools.reduce(Tensor.minimum, data_0)
|
||||
def Sum(*data_0): return functools.reduce(Tensor.add, data_0)
|
||||
def Mean(*data_0): return Sum(*data_0) / len(data_0)
|
||||
# NOTE: does not support saturate
|
||||
def Cast(x: Tensor, to: int, saturate=1): return x.cast(DTYPE_MAP[to])
|
||||
def Cast(x: Tensor, to: int, saturate=1): return x.cast(dtype_parse(to))
|
||||
def CastLike(x: Tensor, target_type: Tensor, saturate=1): return x.cast(target_type.dtype)
|
||||
|
||||
# **************** Simple Ops ****************
|
||||
@@ -91,6 +87,14 @@ def Trilu(x: Tensor, k: Union[Tensor, int]=0, upper=1):
|
||||
k = to_python_const(k) if isinstance(k, Tensor) else 0 # onnx passes k as a tensor int64 with one element, default is 0
|
||||
return x.triu(k) if upper else x.tril(k)
|
||||
|
||||
def Slice(data: Tensor, starts:Tensor, ends:Tensor, axes:Optional[Tensor]=None, steps:Optional[Tensor]=None):
|
||||
if axes is None: axes = list(range(data.ndim))
|
||||
if steps is None: steps = [1] * data.ndim
|
||||
starts, ends, axes, steps = (to_python_const(x) for x in (starts, ends, axes, steps))
|
||||
slices = [slice(0,x,1) for x in data.shape]
|
||||
for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
|
||||
return data[tuple(slices)]
|
||||
|
||||
def Squeeze(data: Tensor, axes):
|
||||
if isinstance(axes, Tensor): axes = to_python_const(axes)
|
||||
axes = [data._resolve_dim(x) for x in axes]
|
||||
@@ -405,7 +409,7 @@ def Compress(inp: Tensor, condition: Tensor, axis=None):
|
||||
return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))]
|
||||
|
||||
def EyeLike(x: Tensor, dtype=None, k=0):
|
||||
ret = Tensor.eye(cast(int, min(x.shape)), dtype=DTYPE_MAP[dtype] if dtype else x.dtype)
|
||||
ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype) if dtype else x.dtype)
|
||||
return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.size(0)-k) for d in x.shape))
|
||||
|
||||
def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)
|
||||
@@ -424,7 +428,7 @@ def DequantizeLinear(x: Tensor, x_scale: Tensor, x_zero_point: Union[Tensor, int
|
||||
def ImageDecoder(encoded_stream: Tensor, pixel_format="RGB"):
|
||||
try: import PIL.Image
|
||||
except ImportError as e: raise ImportError("Pillow must be installed to use the reference implementation of the ImageDecoder operator") from e
|
||||
img = PIL.Image.open(io.BytesIO(to_python_const(encoded_stream, True)))
|
||||
img = PIL.Image.open(io.BytesIO(to_python_const(encoded_stream)))
|
||||
if pixel_format == "BGR": return Tensor(np.array(img))[:, :, ::-1]
|
||||
if pixel_format == "RGB": return Tensor(np.array(img))
|
||||
if pixel_format == "Grayscale": return Tensor(np.array(img.convert("L"))).unsqueeze(-1) # (H, W) to (H, W, 1)
|
||||
|
||||
Reference in New Issue
Block a user