# ***************************************************************
# Copyright (c) 2023 Jittor. All Rights Reserved.
# Maintainers:
# Dun Liang <randonlang@gmail.com>.
# Meng-Hao Guo <guomenghao1997@gmail.com>
#
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
__version__ = '1.3.9.2'
from jittor_utils import lock
with lock.lock_scope():
ori_int = int
ori_float = float
ori_bool = bool
from . import compiler
from .compiler import LOG, has_cuda
from .compiler import compile_custom_ops, compile_custom_op
import jittor_core
import jittor_core as core
from jittor_core import *
from jittor_core.ops import *
from . import compile_extern
from .compile_extern import mkl_ops, mpi, mpi_ops, in_mpi, rank, world_size
if core.get_device_count() == 0:
has_cuda = compile_extern.has_cuda = compiler.has_cuda = False
from .compile_extern import cudnn, curand, cublas, cufft
from .init_cupy import numpy2cupy
from typing import List, Tuple
import contextlib
import numpy as np
from collections import OrderedDict
from collections.abc import Sequence, Mapping
import types
import pickle
import hashlib
import sys, os
import traceback
if "SKEY" in os.environ:
import jittor_utils.student_queue
[文档]
def dfs_to_numpy(x):
"""
将 Jittor 的变量对象转换为 :math:`numpy` 数组的函数。
函数使用深度优先搜索 :math:`（DFS）` 的方式，遍历输入数据中的所有元素，对每个元素进行判断，如果某个元素是 Jittor 的 :math:`var` 对象，就将其转换为 :math:`numpy` 数组。否则，保持不变。
参数：
- x(list, dict, Var 或其它类型)：输入数据，可以是 :math:`Python` 的内置类型（list，dict）或者 Jittor 的 :math:`Var` 对象，或者是这些类型的嵌套。如果输入数据包含 :math:`Var` 对象，它将被转换为 :math:`numpy` 数组。其它类型的数据将保持不变。
返回值：
输入数据的转换结果 (类型和输入变量一致)。所有的 ``Jittor Var`` 对象都被转换为 ``numpy`` 数组，其余数据不变。
代码示例：
>>> import numpy as np
>>> a = jt.ones([2, 3]))
>>> b = [a, a]
>>> c = {\"x\": a, \"y\": b}
>>> res_c = dfs_to_numpy(c)
>>> isinstance(res_c[\"x\"], np.ndarray)
True
>>> isinstance(res_c[\"y\"][0], np.ndarray)
True
"""
if isinstance(x, list):
for i in range(len(x)):
x[i] = dfs_to_numpy(x[i])
elif isinstance(x, dict):
for k in x:
x[k] = dfs_to_numpy(x[k])
elif isinstance(x, Var):
return x.numpy()
return x
[文档]
def safepickle(obj, path):
"""
该函数使用pickle模块将指定对象保存至给定路径。
参数：
- obj(any):需要被保存的对象。
- path(str): 一个字符串，用于指定对象保存的文件路径。
返回值：
无返回值。
代码示例:
>>> import jittor as jt
>>> a = [1, 2, 3]
>>> jt.safepickle(a, '/path/to/save')
"""
if path.endswith(".pth") or path.endswith(".pt") or path.endswith(".bin"):
from jittor_utils.save_pytorch import save_pytorch
save_pytorch(path, obj)
return
# Protocol version 4 was added in Python 3.4. It adds support for very large objects, pickling more kinds of objects, and some data format optimizations.
# ref: <https://docs.python.org/3/library/pickle.html>
# obj = dfs_to_numpy(obj)
s = pickle.dumps(obj, 4)
checksum = hashlib.sha1(s).digest()
s += bytes(checksum)
s += b"HCAJSLHD"
with open(path, 'wb') as f:
f.write(s)
def _load_pkl(s, path):
"""函数用于尝试利用 :math:`pickle` 的 :math:`loads`
函数从指定路径读取 :math:`pickle` 文件。当读取(解码)失败时，
向用户抛出一个具有相关错误信息的 :math:`RuntimeError` 。
参数：
- s(str or bytes): :math:`pickle` 文件流
- path(str): :math:`pickle` 文件的路径
返回值：
:math:`pickle` 文件的解码结果(str)。
"""
try:
return pickle.loads(s)
except Exception as e:
msg = str(e)
msg += f"\nPath: \"{path}\""
if "trunc" in msg:
msg += "\nThis file maybe corrupted, please consider remove it" \
" and re-download."
raise RuntimeError(msg)
def _upload(path, url, jk, tdir=""):
tdir = tdir + '/' if tdir != "" else ""
prefix = f"https://cg.cs.tsinghua.edu.cn/jittor/{tdir}assets"
if url.startswith("jittorhub://"):
url = url.replace("jittorhub://", prefix+"/build/checkpoints/")
assert url.startswith(prefix)
suffix = url[len(prefix):]
dir_suffix = "/".join(suffix.split("/")[:-1])
jkey = flags.cache_path+"/_jkey"
with open(jkey, 'w') as f:
f.write(jk)
assert os.system(f"chmod 600 \"{jkey}\"") == 0
print(dir_suffix)
assert os.system(f"s""s""h"f" -i \"{jkey}\" jittor" "@" "166" f".111.68.30 mkdir -p Documents/jittor-blog/{tdir}assets{dir_suffix}") == 0
assert os.system(f"s""c""p"+f" -i \"{jkey}\" \"{path}\" jittor" "@" "166" f".111.68.30:Documents/jittor-blog/{tdir}assets{suffix}") == 0
assert os.system(f"s""s""h"f" -i \"{jkey}\" jittor" "@" "166" ".111.68.30 Documents/jittor-blog.git/hooks/post-update") == 0
[文档]
def safeunpickle(path):
"""
该函数用于安全地反序列化pickle文件。
参数:
path (str): pickle文件的路径。
返回值:
反序列化的对象(obj)。
示例代码:
>>> import jittor as jt
>>> jt.safeunpickle('path_to_pickle_file')
"""
if path.startswith("jittorhub://"):
path = path.replace("jittorhub://", f"https://cg.cs.tsinghua.edu.cn/jittor/assets/build/checkpoints/")
if path.startswith("https:") or path.startswith("http:"):
base = path.split("/")[-1]
fname = os.path.join(compiler.ck_path, base)
from jittor_utils.misc import download_url_to_local
download_url_to_local(path, base, compiler.ck_path, None)
path = fname
if not (path.endswith(".pth") or path.endswith(".pkl") or path.endswith(".pt")):
return path
if path.endswith(".pth") or path.endswith(".pt") or path.endswith(".bin") :
from jittor_utils.load_pytorch import load_pytorch
model_dict = load_pytorch(path)
return model_dict
with open(path, "rb") as f:
s = f.read()
if not s.endswith(b"HCAJSLHD"):
return _load_pkl(s, path)
checksum = s[-28:-8]
s = s[:-28]
if hashlib.sha1(s).digest() != checksum:
raise ValueError("Pickle checksum does not match! path: "+path,
" This file maybe corrupted, please consider remove it"
" and re-download.")
return _load_pkl(s, path)
class _call_no_record_scope:
"""
此类是一个上下文管理器。重写了 ``__enter__`` 、 ``__exit__`` 和 ``__call__`` 方法，
使得在其管理的上下文中调用的函数仍递归地处于该类作用范围内。使用 ``with`` 语句可将此类用作上下文管理器调用。
其中，``__enter__(self)`` 是在进入上下文环境时调用的方法。``__exit__(self, *exc)`` 是在退出上下文环境时调用的方法。``__call__(self, func)`` 是将该类实例作为装饰器应用到函数上时调用的方法。flag_scope、no_grad 等类是此类重写了这些函数的派生类。
参数：
- func(callable)：可调用的函数。
- args : 位置参数。
- kw : 关键字参数。
- exc：用于接收可能发生的异常信息的可变参数。
代码示例：
>>> @jt._call_no_record_scope()
... def my_function():
... print(\"Inside my_function\")
>>> my_function()
Inside my_function
"""
def __enter__(self): pass
def __exit__(self, *exc): pass
def __call__(self, func):
def inner(*args, **kw):
with self:
ret = func(*args, **kw)
return ret
return inner
[文档]
class flag_scope(_call_no_record_scope):
"""
该类用于构建 Jittor 中标志作用域，及将 Jittor 配置或标志在这个作用域内进行变动。该类只是一个上下文管理器类，用于临时设置标志位的值，并在离开作用域时恢复为原始值。
"""
def __init__(self, **jt_flags):
self.jt_flags = jt_flags
def __enter__(self):
flags_bk = self.flags_bk = {}
try:
for k,v in self.jt_flags.items():
origin = getattr(flags, k)
flags_bk[k] = origin
# merge dict attrs
if isinstance(origin, dict):
for ok, ov in origin.items():
if ok not in v:
v[ok] = ov
setattr(flags, k, v)
except:
self.__exit__()
raise
def __exit__(self, *exc):
for k,v in self.flags_bk.items():
setattr(flags, k, v)
[文档]
class no_grad(flag_scope):
""" 取消梯度范围，所有在此范围内创建的变量都将停止梯度。
继承自 flag_scope 类，作为上下文管理器，通过将 jt_flags 的 no_grad 键值设为 1，使优化器取消计算其范围内每个运算的梯度。
参数:
jt_flags (dict): 包含若干 Jittor 标记的字典。默认值：空。
代码示例：
>>> import jittor as jt
>>> x = jt.array([1, 2, 3], dtype=float)
>>> y = jt.array([4, 5, 6], dtype=float)
>>> with jt.no_grad():
... z = x * y
... w = z.sum()
>>> optimizer = jt.optim.SGD(params=[x, y, z, w], lr=0.1)
>>> optimizer.step(w)
>>> x.opt_grad(optimizer)
jt.Var([0. 0. 0.], dtype=float64)
"""
def __init__(self, **jt_flags):
self.jt_flags = jt_flags
jt_flags["no_grad"] = 1
[文档]
class enable_grad(flag_scope):
'''
在scope类开启梯度。所有在这个scope内创建的变量都会记录梯度。
示例代码::
import jittor as jt
with jt.enable_grad():
...
'''
def __init__(self, **jt_flags):
self.jt_flags = jt_flags
jt_flags["no_grad"] = 0
single_log_capture = None
[文档]
class log_capture_scope(_call_no_record_scope):
"""
该类是一个上下文管理器，用于捕获并记录 log 信息。所有在该类的上下文管理器内的日志记录信息都将被捕获并保存。该类继承自 _call_no_record_scope 类，通过在代码段中临时地改变 jt_flags 标志来实现记录 LOG 调用信息。
参数：
- \\*\\*jt_flags (dict):用于修改 Jittor 运行时的参数，接受任意数量的键值对作为输入。
属性：
- fs: 一个 flag_scope 对象，用于备份和修改 Jittor 运行时的参数。
- logs: 一个空列表，在进入 log_capture_scope 时被创建。这个列表用于存储捕获到的日志信息。
代码示例：
>>> with jt.log_capture_scope(log_v=1) as logs:
... jt.LOG.v(\"...\")
>>> print(logs)
[{'level': 'i', 'lineno': '2', 'msg': '...', 'name': '<stdin>', 'verbose': '1'}]
"""
def __init__(self, **jt_flags):
jt_flags["use_parallel_op_compiler"] = 0
self.fs = flag_scope(**jt_flags)
def __enter__(self):
global single_log_capture
assert not single_log_capture
single_log_capture = True
self.logs = []
LOG.log_capture_start()
try:
self.fs.__enter__()
if "log_v" in self.fs.jt_flags:
LOG.log_v = self.fs.jt_flags["log_v"]
return self.logs
except:
LOG.log_capture_stop()
single_log_capture = None
raise
def __exit__(self, *exc):
global single_log_capture
self.fs.__exit__(*exc)
if "log_v" in self.fs.jt_flags:
LOG.log_v = flags.log_v
LOG.log_capture_stop()
self.logs.extend(LOG.log_capture_read())
single_log_capture = None
[文档]
class profile_scope(_call_no_record_scope):
"""
profile_scope 类是用来测量代码执行时间的一个工具，封装了自动化性能分析的功能。使用此上下文管理器来记录 Jittor 程序的性能数据。
"""
def __init__(self, warmup=0, rerun=0, **jt_flags):
self.fs = flag_scope(**jt_flags)
self.warmup = warmup
self.rerun = rerun
def __enter__(self):
assert not flags.profiler_enable
self.report = []
try:
self.fs.__enter__()
profiler.start(self.warmup, self.rerun)
return self.report
except:
profiler.stop()
raise
def __exit__(self, *exc):
profiler.stop()
self.report.extend(profiler.report())
self.fs.__exit__(*exc)
[文档]
class profile_mark(_call_no_record_scope):
"""
profile_mark 类用于分析部分代码的性能。该类继承了 _call_no_record_scope 类，作为上下文管理器，通过临时改变全局变量 flags.compile_options 来实现统计其范围内代码的运行时间和内存等性能参数。
代码示例：
>>> a = jt.rand(1000,1000)
>>> b = jt.rand(1000,1000)
>>> jt.sync_all()
>>> results = []
>>> with jt.profile_scope() as rep:
... results.append(jt.matmul(a, b))
... with jt.profile_mark(\"mark1\"):
... results.append(jt.matmul(a, b))
... with jt.profile_mark(\"mark2\"):
... results.append(jt.matmul(a, b))
... with jt.profile_mark(\"mark3\"):
... results.append(jt.matmul(a, b))
... results.append(jt.matmul(a, b))
# 输出：性能分析结果，包括每个 mark 段落的性能评估报告
"""
def __init__(self, mark_name: str):
''' profiler mark is used for profiling part of code,
Example::
a = jt.rand(1000,1000)
b = jt.rand(1000,1000)
jt.sync_all()
results = []
with jt.profile_scope() as rep:
results.append(jt.matmul(a, b))
with jt.profile_mark("mark1"):
results.append(jt.matmul(a, b))
with jt.profile_mark("mark2"):
results.append(jt.matmul(a, b))
with jt.profile_mark("mark3"):
results.append(jt.matmul(a, b))
results.append(jt.matmul(a, b))
Output::
Total time: 46.8ms
Total Memory Access: 57.2MB
[Mark mark3] time: 9ms
[Mark mark2] time: 8.28ms
[Mark mark1] time: 17.7ms
'''
self.mark_name = mark_name
def __enter__(self):
self.options = flags.compile_options
new_options = flags.compile_options
prev_marks = "_marks:"
for x in self.options:
if x.startswith(prev_marks):
prev_marks = x
del new_options[x]
new_marks = prev_marks + self.mark_name + ','
new_options[new_marks] = 1
flags.compile_options = new_options
def __exit__(self, *exc):
flags.compile_options = self.options
class __single_process_scope:
"""
此类作为上下文管理器，用于控制进程范围在单一进程（即主进程）内。在使用并行化工具如 MPI 时，有时需要控制代码仅在主进程中运行。此时可以调用此类进行范围限制。使用 ``with`` 语句调用时，其中的代码块将根据MPI世界排名与给定排名进行比较，若相等，则在这个进程中执行，否则跳过。
当 Jittor 在 MPI 环境下（检查 ``jt.in_mpi`` 为 ``True``）时，此类范围内代码将不使用 MPI 而转为单进程执行，离开范围后恢复原 MPI 环境。
参数：
rank (int, optional)：指定进程的排名。默认值：0
代码示例：
>>> with jt.__single_process_scope(rank=0):
... print(\"I am rank 0\")
>>> print(\"Hello, world!\")
"""
def __init__(self, rank=0):
self.rank = rank
def __enter__(self):
global in_mpi
self.bk_in_mpi = in_mpi
if mpi:
self.bk_mpi_state = mpi.get_state()
if not in_mpi:
return True
ret = self.rank == mpi.world_rank()
in_mpi = compile_extern.in_mpi = False
mpi.set_state(False)
return ret
def __exit__(self, *exc):
global in_mpi
in_mpi = compile_extern.in_mpi = self.bk_in_mpi
if mpi:
mpi.set_state(self.bk_mpi_state)
[文档]
def single_process_scope(rank=0):
"""
为指定的单个进程建立独立的作用域。
参数:
- rank (int): 对应进程的的编号。
返回值:
无返回值
代码示例:
>>> @jt.single_process_scope(rank=0)
>>> def xxx():
>>> ...
"""
def outer(func):
def inner(*args, **kw):
"""
该函数接收任意数量的位置参数和关键字参数，并将他们传递给自身作为函数调用。
参数：
\\*args (任意类型) : 任意数量的位置参数
\\**kw (任意类型) : 任意数量的关键字参数
返回：
如果函数执行成功，返回函数的结果。否则，将返回None。
"""
ret = None
sync_all()
with __single_process_scope(rank) as flag:
if flag:
ret = func(*args, **kw)
return ret
return inner
return outer
[文档]
def clean():
"""
清理和重置Jittor的内部状态。
参数:
不需要任何参数。
返回值:
无。
代码示例:
>>> import jittor as jt
>>> jt.clean()
"""
import gc
# make sure python do a full collection
gc.collect()
core.gc()
# cast = unary
# Var.cast = Var.cast
[文档]
def array(data, dtype=None):
"""
将输入的对象转化为 ``Jittor`` 的Var对象。 即函数主要用于将给定的对象转化为 ``Jittor`` 框架可以处理的 ``Var`` 对象。如果输入的是数值，则创建的 ``Var`` 的 ``dtype`` 为输入的数值类型，若是``numpy.ndarray`` ，则 ``Var`` 的 ``dtype`` 以 ``ndarray`` 的 ``dtype`` 为准；若输入为 ``Var`` 对象，其 ``dtype`` 保持不变。
参数：
- obj(float | int | numpy.ndarray | Var): 需要转化的对象。
返回值：
对应的 Var 对象
代码示例：
>>> import jittor as jt
>>> import numpy as np
>>> jt.array(1)
Var(1, dtype=int32)
>>> jt.array(1.0)
Var(1.0, dtype=float32)
>>> x=jt.array(np.array([1,2,3]))
>>> print(x)
Var([1 2 3], dtype=int32)
>>> jt.array(x)
Var([1 2 3], dtype=int32)
"""
if isinstance(data, core.Var):
if dtype is None:
ret = data.clone()
else:
ret = cast(data, dtype)
elif dtype is not None:
if isinstance(dtype, NanoString):
dtype = str(dtype)
elif callable(dtype):
dtype = dtype.__name__
with jt.flag_scope(auto_convert_64_to_32=0):
ret = ops.array(np.array(data, dtype))
else:
ret = ops.array(data)
# TODO: move those code to core
amp_reg = jt.flags.amp_reg
if amp_reg and ret.numel() != 1 and ret.dtype.is_float():
if amp_reg & 16:
if amp_reg & 1:
if ret.dtype != "float32":
return ret.float32()
elif amp_reg & 2:
if ret.dtype != "float16":
return ret.float16()
return ret
[文档]
def random(shape, dtype="float32", type="uniform"):
"""
生成一个给定形状，值为随机数的张量
参数：
- shape (tuple[int]): 生成张量的形状
- dtype (str): 数据类型。默认值： ``'float32'``
- type (Literal['uniform', 'normal']): ``'uniform'`` ：生成 0 到 1 之间的均匀随机数； ``'normal'`` ：生成均值为 0 方差为 1 的正态分布随机数。默认值： ``'uniform'``
返回值：
生成的随机数
代码示例：
>>> jt.random((3,))
jt.Var([0.7427739 0.80093205 0.2652795 ], dtype=float32)
>>> jt.random((3,), type='normal')
jt.Var([ 1.2703565 0.20411628 -1.1991016 ], dtype=float32)
"""
# TODO: move those code to core
if dtype in ["float16", "bfloat16"]:
# TODO: make curand support fp16
ret = ops.random(shape, "float32", type).cast(dtype)
else:
ret = ops.random(shape, dtype, type)
amp_reg = jt.flags.amp_reg
if amp_reg:
if amp_reg & 16:
if amp_reg & 1:
if ret.dtype != "float32":
return ret.float32()
elif amp_reg & 2:
if ret.dtype != "float16":
return ret.float16()
return ret
[文档]
def float_auto(x):
"""
将输入的参数自动转化为浮点数。
参数:
- x (int, float, list, tuple, np.ndarray) : 输入的待转化的参数，可以是数值型数据或者是类似列表、数组这样的容器型数据。容器型数据中的每一个元素都会被转化为相应的浮点数。
返回值:
转化后的浮点数。如果输入的是容器型数据，则返回的也是对应的容器型数据。比如输入的是列表，则返回的是浮点数列表；输入的是np.ndarray，则返回的是浮点数np.ndarray。
代码示例:
>>> jt.float_auto(1)
1.0
>>> jt.float_auto([1, 2, 3])
[1.0, 2.0, 3.0]
"""
if jt.flags.amp_reg & 2:
return x.float16()
return x.float32()
Var.float_auto = float_auto
[文档]
def array64(data, dtype=None):
'''
本函数将输入转换为 ``int64`` 类型的张量。
参数:
- data(array_like): 任何可以被转换为 ``jittor.Var`` 的数据形式，可以是list，list的list，元组等。
- dtype(Any | None, 可选): 输出的数据的数据类型，当缺省时，生成的数组将具有 ``numpy.float64`` 的数据类型。默认值: ``None``
返回值:
jittor.Var: 一个 ``int64`` 类型的张量
代码示例：
>>> import jittor as jt
>>> jt.array64([1,2,3])
jt.Var([1 2 3], dtype=int32)
'''
with jt.flag_scope(auto_convert_64_to_32=0):
return array(data, dtype)
[文档]
def grad(loss, targets, retain_graph=True):
"""
计算损失函数对目标变量的梯度。对于每个目标变量，求取损失函数对该变量的偏导数，这个偏导数就是该变量的梯度。
参数:
- loss (Var) : 损失函数值。
- targets (Var): 目标变量，可以是一个变量或变量的列表。
- retain_graph (bool) : 是否保持计算图。默认值为False。如果设为True，那么在计算梯度之后，计算图不会被删除，可以进行多次反向传播。
返回值:
返回计算得到的梯度张量(Var)。
"""
if type(targets) == core.Var:
return core.grad(loss, [targets], retain_graph)[0]
return core.grad(loss, targets, retain_graph)
[文档]
def liveness_info():
"""
获取Jittor库中的存活信息。
参数:
无
返回值:
包含存活信息的字典(dict)，包括 ``hold_vars`` , ``lived_vars`` , ``lived_ops`` 。
代码示例：
>>> jt.liveness_info()
{'hold_vars': 0, 'lived_vars': 0, 'lived_ops': 0}
"""
return {
"hold_vars": core.number_of_hold_vars(),
"lived_vars": core.number_of_lived_vars(),
"lived_ops": core.number_of_lived_ops(),
}
[文档]
def ones(*shape, dtype="float32"):
"""
创建一个指定形状和数据类型的全 1 矩阵。
参数:
- shape (int or sequence of ints): 形成全 1 矩阵的形状。可以是单个整数（生成一维数组）或整数序列（生成多维数组）。
- dtype (str, 可选): 矩阵的数据类型，默认为 ``float``。
返回值:
- 返回一个给定形状和数据类型的全 1 矩阵，类型为 jt.Var。
代码示例:
>>> import jittor as jt
>>> jt.ones(5)
jt.Var([1. 1. 1. 1. 1.], dtype=float32)
>>> jt.ones((5,), dtype='int')
jt.Var([1 1 1 1 1], dtype=int32)
>>> jt.ones((2, 1), dtype='float')
jt.Var([[1.]
[1.]], dtype=float32)
"""
if isinstance(shape, tuple) and isinstance(shape[-1], (str, NanoString)):
dtype = shape[-1]
shape = shape[:-1]
if isinstance(shape, tuple) and isinstance(shape[0], (Sequence, NanoVector)):
shape = shape[0]
return unary(1, dtype).broadcast(shape)
[文档]
def new_ones(x, size):
"""
返回一个全为1的张量(Var)，形状和元素类型与输入张量 ``x`` 定义保持一致。
参数：
- x (Var): Var类型的张量
返回值：
全为1的张量(Var)且 ``dtype`` 和 ``x`` 保持一致
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.new_ones(x, (5,))
jt.Var([1. 1. 1. 1. 1.], dtype=float32)
"""
return ones(size, x.dtype)
Var.new_ones = new_ones
[文档]
def ones_like(x):
"""
返回一个全为1的张量(Var)，形状和元素类型与输入张量 ``x`` 保持一致。
参数：
- x (Var): Var类型的张量
返回值：
全为1的张量(Var)，其形状和元素类型与输入张量 ``x`` 保持一致
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.ones_like(x)
jt.Var([1. 1. 1.], dtype=float32)
"""
return ones(x.shape,x.dtype)
[文档]
def zeros(*shape, dtype="float32"):
"""
返回一个全为0的张量(Var)，形状由可变参数 ``size`` 定义，数据类型由 ``dtype`` 定义。如果不给定 ``dtype`` , 默认类型为 ``float32``。
参数：
- size (int...): 整数序列，定义了输出的形状
- dtype (var.dtype, optional): 数据类型，默认为 ``float32``
返回值：
元素全为0的张量(Var)
代码示例：
>>> jt.zeros((3,2), dtype='int32')
jt.Var([[0 0]
[0 0]
[0 0]], dtype=int32)
>>> jt.ones_like(x)
jt.Var([1. 1. 1.], dtype=float32)
"""
if isinstance(shape, tuple) and isinstance(shape[-1], (str, NanoString)):
dtype = shape[-1]
shape = shape[:-1]
if isinstance(shape, tuple) and isinstance(shape[0], (Sequence, NanoVector)):
shape = shape[0]
return unary(0, dtype).broadcast(shape)
[文档]
def new_zeros(x, size):
"""
返回一个全为0的张量(Var)，数据类型由张量 ``x`` 定义，形状由可变参数 ``size`` 定义。
参数：
- x (Var): Var类型的张量
- size (int...): 整数序列，定义了输出的形状
返回值：
全为0的张量(Var)且数据类型 ``dtype`` 和输入张量 ``x`` 保持一致
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.new_zeros(x, (5,))
jt.Var([0. 0. 0. 0. 0.], dtype=float32)
"""
return zeros(size, x.dtype)
Var.new_zeros = new_zeros
[文档]
def empty(*shape, dtype="float32"):
"""
通过指定的形状和数据类型创建一个空的张量。
参数：
- shape(Tuple[int]):创建的张量的形状。
- dtype(str,可选):创建的张量的数据类型，默认值：``float32``。
返回值：
创建的空的张量(Var)
代码示例：
>>> import jittor as jt
>>> jt.empty((2, 3))
"""
if isinstance(shape, tuple) and isinstance(shape[-1], (str, NanoString)):
dtype = shape[-1]
shape = shape[:-1]
if isinstance(shape, tuple) and isinstance(shape[0], (Sequence, NanoVector)):
shape = shape[0]
return ops.empty(shape, dtype)
[文档]
def new_empty(x, size):
"""
返回一个用未初始化数据填充的张量(Var)，形状和元素类型与输入张量 ``x`` 定义保持一致。
返回值：
返回一个用初始化数据填充的张量(Var)且数据类型 ``dtype`` 和输入张量 ``x`` 保持一致
参数：
x (Var): Var类型的张量
size (int...): 整数序列，定义了输出的形状
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.new_empty(x, (5,))
jt.Var([0. 0. 0. 0. 0.], dtype=float32)
"""
return empty(size, x.dtype)
Var.new_empty = new_empty
[文档]
def full(shape,val,dtype="float32"):
"""
创建一个张量(Var)，其中张量的形状由可变参数 ``shape`` 定义，
张量中的各个标量的值由可变参数 ``val`` 定义，数据类型由 ``dtype`` 定义。如果不给定 ``dtype`` , 默认类型为 ``float32``。
参数：
shape (int...): 整数序列，定义了输出张量的形状
val: 填充值，任意数值类型
dtype (var.dtype, optional): 数据类型，默认为 ``float32``
返回值：
返回一个填充值为 ``val`` ，形状大小为 ``shape`` 的张量(Var)
代码示例：
>>> jt.full((3,2), 3.14)
jt.Var([[3.14 3.14]
[3.14 3.14]
[3.14 3.14]], dtype=float32)
"""
if not isinstance(shape, (NanoVector, Sequence)):
shape = (shape,)
return unary(val, dtype).broadcast(shape)
[文档]
def new_full(x, size, val):
"""
创建一个张量(Var)，其中张量的形状由可变参数 ``shape`` 定义，张量中的各个标量的值由可变参数 ``val`` 定义，数据类型由张量 ``x`` 的数据类型定义。
参数：
- x (Var): Var类型的张量
- size(int...): 整数序列，定义了输出张量的形状
- val: 填充值，任意数值类型
返回值：
返回一个填充值为 ``val`` ，形状大小为 ``size`` 的张量(Var)。
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.new_full(x, (3,2), 3)
jt.Var([[3. 3.]
[3. 3.]
[3. 3.]], dtype=float32)
"""
return full(size, val, x.dtype)
Var.new_full = new_full
def ne(x,y):
"""
将两个输入进行比较并返回结果。如果 ``x`` 与 ``y`` 不相等，返回True；否则返回False。
``x`` 和 ``y`` 必须是可以被广播为相同形状的数组，满足以下条件的数组可以执行广播操作。
- 相同形状的数组可以进行广播。
- 任何维度为1的数组都可以广播。
- 在某个维度上大小不匹配的数组不能广播，除非其中一个数组在该维度的大小为1。
参数：
- x(int, float, list, numpy.array):输入的第一个数值，可以是数字、列表或者 :math:`numpy` 数组
- y(int, float, list, numpy.array): 输入的第二个数值，可以是数字、列表或者 :math:`numpy` 数组。注意 ``y`` 必须与 ``x`` 相同的维度和大小
返回值：
与 ``x`` 和 ``y`` 相同形状和维度的布尔值数组，其中的元素表示 ``x`` 和 ``y`` 对应位置元素的不相等性。如果 ``x`` != ``y`` , 结果为True，否则为False
代码示例：
>>> x = [1, 2, 3]
>>> y = [1, 2, 4]
>>> ne(x, y)
array([False, False, True])
"""
return x!=y
Var.ne = ne
[文档]
def full_like(x, val, dtype=None) -> Var:
"""
创建一个张量(Var)，其中张量的形状由张量 ``x`` 的形状定义，张量中的各个标量的值由可变参数 ``val`` 定义，数据类型默认为 ``None`` ，即张量 ``x`` 的数据类型。
参数：
- x (Var): Var类型的张量
- val: 填充值，任意数值类型
- dtype (var.dtype, optional): 数据类型，默认为None
返回值：
返回一个填充值为 ``val`` ，形状大小为和 ``x`` 一致的张量(Var)
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.full_like(x, 3.14)
jt.Var([3.14 3.14 3.14], dtype=float32)
"""
if dtype is None: dtype = x.dtype
return full(x.shape, val, dtype)
[文档]
def zeros_like(x, dtype=None) -> Var:
"""
创建一个全为0的张量(Var)，其中张量的形状由张量 ``x`` 的形状定义，数据类型默认为 ``None`` ,即张量 ``x`` 的数据类型。
参数：
- x (Var): Var类型的张量
- dtype (var.dtype, optional): 数据类型，默认为 ``None``
返回值：
返回一个值全为0，形状大小为和 ``x`` 一致的张量(Var)
代码示例：
>>> jt.array([1, 2, 3]).float()
jt.Var([1. 2. 3.], dtype=float32)
>>> jt.zeros_like(x)
jt.Var([0. 0. 0.], dtype=float32)
"""
if dtype is None: dtype = x.dtype
return zeros(x.shape, dtype)
flags = core.Flags()
[文档]
def var(x, dim=None, dims=None, unbiased=False, keepdims=False):
"""
返回样本对指定维度的方差。如果``unbiased``为True，则使用贝塞尔修正。
参数:
- x (jt.Var) : 输入的jittor变量。
- dim (int) : 计算方差的维度。如果dim和dims都为None，则将计算整个张量的方差。默认值: ``None``
- dims (tuple of int) : 计算方差的维度。如果dim和dims都为None，则将计算整个张量的方差。默认值: ``None``
- unbiased (bool) : 如果为True，则使用贝塞尔修正。
- keepdims (bool) : 用于指定计算的输出是否保持原始张量的维度。如果 ``keepdims=True``，输出形状与输入形状相同。比如，输入的张量形状是 (3, 4) ，方差是对第一维（``dim=0``）求的，则输出的形状还是 (3, 4)。如果 ``keepdims=False``，则输出的形状就是 (4,)。默认值: ``False`` 。
返回值:
返回计算得到的方差,类型为jittor.Var。
代码示例：
>>> a = jt.rand(3)
>>> a
jt.Var([0.79613626 0.29322362 0.19785859], dtype=float32)
>>> a.var()
jt.Var([0.06888353], dtype=float32)
>>> a.var(unbiased=True)
jt.Var([0.10332529], dtype=float32)
"""
shape = x.shape
new_shape = list(x.shape)
assert dim is None or dims is None, "dim and dims can not be both set"
if dim is None and dims is None:
dims = list(range(len(shape)))
elif dim is not None:
dims = [dim]
mean = jt.mean(x, dims, keepdims=True)
mean = jt.broadcast(mean, shape)
n = 1
for d in dims:
n *= shape[d]
new_shape[d] = 1
sqr = (x - mean) ** 2
sqr = jt.sum(sqr, dims=dims, keepdims=False)
if unbiased:
n -= 1
sqr /= n
if keepdims:
sqr = sqr.view(new_shape)
return sqr
Var.var = var
[文档]
def std(x):
"""
计算给定张量 ``x`` 的标准差。
标准差的计算公式如下:
.. math::
\\sigma=\\sqrt{\\sum_{i=0}^{N-1}\\left(x_i-\\bar{x}\\right)^2}
其中 :math:`\\bar{x}` 为 ``x`` 的均值, :math:`N` 为 ``x`` 中元素的数量， :math:`x_i` 为 ``x`` 中的第i个元素
参数:
- x (jt.Var) : 输入的jittor张量。
返回值:
输入张量 ``x`` 的标准差,类型为jittor.Var。
代码示例：
>>> x = jt.array([1.0, 2.0, 3.0, 4.0])
>>> jt.std(x)
jt.Var([1.2909944], dtype=float32)
"""
matsize=1
for i in x.shape:
matsize *= i
out=(x-x.mean()).sqr().sum()
out=out/(matsize-1)
out=out.maximum(1e-6).sqrt()
return out
Var.std = std
[文档]
def norm(x, p=2, dim=-1, keepdims=False, eps=1e-30, keepdim=False):
"""
计算输入向量 ``x`` 的 ``p`` 范数(默认为2范数), ``p`` 只能为1或2。
如果p=1，计算公式如下:
.. math::
\\|x\\|_1=\\left|x_1\\right|+\\left|x_2\\right|+\\ldots+\\left|x_n\\right|
如果p=2，计算公式如下：
.. math::
\\|x\\|_2=\\left(\\left|x_1\\right|^2+\\left|x_2\\right|^2+\\ldots+\\left|x_n\\right|^2\\right)^{1 / 2}
参数:
- x (jt.Var) : 输入的jittor张量。
- p (int) : 需要计算的范数的类型，只能为1或2，默认值:2。
- dim (int) : 指定需要计算范数的维度，如果为-1,则计算整个张量的范数。默认值:-1。
- keepdims (bool) : 指定是否在输出的结果中保留原始的维度。默认值:False。
- eps (float) : 用于保持数值稳定性的微小值。默认值:1e-12。
- keepdim (bool) : 指定是否在输出的结果中保留原始的维度。默认值:False。
返回值:
计算的p范数的结果,类型为jittor.Var。
代码示例：
>>> x = jt.array([1.0, 2.0, 3.0])
>>> jt.norm(x)
jt.Var([3.7416575], dtype=float32)
"""
keepdim = keepdim or keepdims
assert p==1 or p==2
if p==1:
return x.abs().sum(dim, keepdim)
if p==2:
return (x.sqr()).sum(dim, keepdim).maximum(eps).sqrt()
Var.norm = norm
origin_reshape = reshape
[文档]
def reshape(x, *shape):
"""
将输入的张量根据指定的形状进行重塑，并返回这个新的张量。保持元素的数据和数量不变。其中一维可以是-1，在这种情况下，它是根据其余的维度和输入的元素数量推断出来的。
参数：
- x(jt.Var): 输入的张量。
- shape(Tuple[int]): 输出的形状,一个整数列表。
返回值：
重塑过后的张量 （jt.Var）
代码示例：
>>> jt.randint(0, 10, shape=(12,))
jt.Var([4 0 8 4 6 3 1 8 1 1 2 2], dtype=int32)
>>> jt.reshape(a, (3, 4))
jt.Var([[4 0 8 4]
[6 3 1 8]
[1 1 2 2]], dtype=int32)
>>> jt.reshape(a, (-1, 6))
jt.Var([[4 0 8 4 6 3]
[1 8 1 1 2 2]], dtype=int32)
"""
if len(shape) == 1 and isinstance(shape[0], (Sequence, NanoVector)):
shape = shape[0]
return origin_reshape(x, shape)
reshape.__doc__ = origin_reshape.__doc__
Var.view = Var.reshape = view = reshape
origin_transpose = transpose
[文档]
def transpose(x, *dim):
"""
返回一个张量，其数据与原始张量 ``x`` 相同，其指定的维度已经根据 ``dim`` 参数中的指示进行了交换。
参数：
- x (Var): Var类型的张量
- dim (int...): 整数序列，表示需要交换的维度, 默认为空，表示对所有轴进行转置
代码示例：
>>> x = jt.randn((3,4))
>>> x
jt.Var([[ 1.7929314 0.8375565 0.13141492 1.4283884 ]
[ 0.0627789 1.4741095 -0.3331003 0.34370992]
[ 0.405448 1.6529875 0.19640142 -0.09192007]], dtype=float32)
>>> jt.transpose(x)
jt.Var([[ 1.7929314 0.0627789 0.405448 ]
[ 0.8375565 1.4741095 1.6529875 ]
[ 0.13141492 -0.3331003 0.19640142]
[ 1.4283884 0.34370992 -0.09192007]], dtype=float32)
返回值：
返回改变维度的张量（Var）
"""
if len(dim) == 1 and isinstance(dim[0], (Sequence, NanoVector)):
dim = dim[0]
elif len(dim) == 2:
axes = list(range(x.ndim))
a, b = dim
axes[a], axes[b] = axes[b], axes[a]
dim = axes
return origin_transpose(x, dim)
transpose.__doc__ = origin_transpose.__doc__
Var.transpose = Var.permute = transpose
[文档]
def permute(x, *dim):
"""
返回一个张量，其数据与原始张量 ``x`` 相同，其指定的维度已经根据 ``dim`` 参数中的指示进行了交换。
参数：
- x (Var): Var类型的张量
- dim (int...): 整数序列，表示需要交换的维度, 默认为空，表示对所有轴进行转置
返回值：
返回改变维度的张量（Var）
代码示例：
>>> x = jt.randn((3,4))
>>> x
jt.Var([[ 1.7929314 0.8375565 0.13141492 1.4283884 ]
[ 0.0627789 1.4741095 -0.3331003 0.34370992]
[ 0.405448 1.6529875 0.19640142 -0.09192007]], dtype=float32)
>>> x.permute(1,0)
jt.Var([[ 1.7929314 0.0627789 0.405448 ]
[ 0.8375565 1.4741095 1.6529875 ]
[ 0.13141492 -0.3331003 0.19640142]
[ 1.4283884 0.34370992 -0.09192007]], dtype=float32)
"""
if len(dim) == 1 and isinstance(dim[0], (Sequence, NanoVector)):
dim = dim[0]
elif len(dim) == 2:
axes = list(range(x.ndim))
a, b = dim
axes[a], axes[b] = axes[b], axes[a]
dim = axes
return origin_transpose(x, dim)
[文档]
def flatten(input, start_dim=0, end_dim=-1):
"""
将输入张量 ``input`` 在指定维度的范围 ``[start_dim, end_dim]`` 上进行展平操作。
参数:
- input (jt.Var) : 输入的jittor张量。
- start_dim (int) : 展平操作的起始维度，默认值:0。
- end_dim (int) : 展平操作的结束维度，默认值:-1(代表最后一个维度)
返回值:
经过展平操作后的新张量,类型为jittor.Var。
代码示例：
>>> x = jt.array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])
>>> y = jt.flatten(x)
jt.Var([1 2 3 4 5 6 7 8], dtype=int32)
>>> y = jt.flatten(x,start_dim=1)
jt.Var([[1 2 3 4]
[5 6 7 8]], dtype=int32)
"""
in_shape = input.shape
start_dim = len(in_shape) + start_dim if start_dim < 0 else start_dim
end_dim = len(in_shape) + end_dim if end_dim < 0 else end_dim
assert end_dim >= start_dim, "end_dim should be larger than or equal to start_dim for flatten function"
out_shape = []
for i in range(0,start_dim,1): out_shape.append(in_shape[i])
dims = 1
for i in range(start_dim, end_dim+1, 1): dims *= in_shape[i]
out_shape.append(dims)
for i in range(end_dim+1,len(in_shape),1): out_shape.append(in_shape[i])
return input.reshape(out_shape)
Var.flatten = flatten
Var.detach_inplace = Var.start_grad
[文档]
def detach(x):
"""
创建一个已经存在的张量 ``x`` 的副本，这个副本与原始张量共享数据，但不参与梯度计算。
参数：
- x (Var): Var类型的张量
返回值：
返回副本张量（Var）
代码示例：
>>> x = jt.randn((3,2))
jt.Var([[-1.7332076 -3.9374337]
[ 2.2316337 1.2128713]
[-2.067835 0.8186041]], dtype=float32)
>>> x.detach()
jt.Var([[-1.7332076 -3.9374337]
[ 2.2316337 1.2128713]
[-2.067835 0.8186041]], dtype=float32)
"""
return x.detach()
[文档]
def unsqueeze(x, dim):
"""
创建一个新张量，其数据与原始张量 ``x`` 相同，其形状在指定的 ``dim`` 维度上增加了一个大小为 1 的新维度。
参数：
- x (Var): Var类型的张量
- dim (int): 插入单维度的索引值，该值的范围需要满足 ``[-x.dim() - 1, x.dim() + 1)`` 。
代码示例：
>>> x = jt.array([1,2,3,4])
jt.Var([1 2 3 4], dtype=int32)
>>> jt.unsqueeze(x, 0)
jt.Var([[1 2 3 4]], dtype=int32)
>>> jt.unsqueeze(x, 1)
jt.Var([[1]
[2]
[3]
[4]], dtype=int32)
返回值：
返回一个新的张量（Var），其形状在指定的 ``dim`` 维度上增加了一个大小为 1 的新维度。其他维度保持不变。
"""
shape = list(x.shape)
if dim < 0: dim += len(shape) + 1
assert dim <= len(shape)
return x.reshape(shape[:dim] + [1] + shape[dim:])
Var.unsqueeze = unsqueeze
[文档]
def squeeze(x, dim=None):
"""
创建一个新张量，移除 输入变量``x`` 张量中所有长度为1的维度。如果指定维度 ``dim``, 则只移除该维度（如果其长度为1）。
参数：
- x (Var): Var类型的张量
- dim (int, optional): 需要删除的维度，默认为 ``None`` 表示所有维度。
返回值：
返回一个新的张量（Var），移除指定维度长度为1的维度。
代码示例：
>>> x = jt.zeros(2, 1, 2, 1, 2)
>>> x.size()
[2,1,2,1,2,]
>>> y1 = jt.squeeze(x)
>>> y1.size()
[2,2,2,]
>>> y2 = jt.squeeze(x, 1)
>>> y2.size()
[2,2,1,2,]
注意事项:
- 如果 ``dim`` 传入了长度不为1的维度，或者多个维度的元组，会引起异常
"""
shape = list(x.shape)
if dim is None:
new_shape = [s for s in shape if s > 1]
return x.reshape(new_shape)
else:
if dim < 0: dim += len(shape)
assert dim < len(shape) and dim >= 0
assert shape[dim] == 1
return x.reshape(shape[:dim] + shape[dim+1:])
Var.squeeze = squeeze
[文档]
def clamp(x, min_v=None, max_v=None):
"""
创建一个新张量，限制张量 ``x`` 中数值在特定的范围 ``[min_v, max_v]`` 内。对于张量 ``x`` 中的任何一个数值均满足：
.. math::
y_i = \\min(\\max(x_i, \\text{min_value}_i), \\text{max_value}_i)
参数：
- x (Var): Var类型的张量
- min_v (int, optional): 数值的最小允许值。如果为 None，则不应用最小值限制
- max_v (int, optional): 数值的最大允许值。如果为 None，则不应用最大值限制
返回值：
返回一个新的张量（Var），其数值的维度被限制在 ``[min_v, max_v]`` 之间 。
代码示例：
>>> x = jt.randn(4)
jt.Var([ 0.8474737 0.9300491 2.8043647 -0.32242754], dtype=float32)
>>> jt.clamp(x, -0.5, 0.5)
jt.Var([ 0.5 0.5 0.5 -0.32242754], dtype=float32)
注意事项:
- 如果 ``min_v`` 大于 ``max_v`` ，会引起异常
"""
if x.shape[0]==0:
return x
if min_v is not None and max_v is not None:
assert min_v <= max_v
if min_v is not None:
x = x.maximum(min_v)
if max_v is not None:
x = x.minimum(max_v)
return x
Var.clamp = clamp
[文档]
def clamp_(x, min_v=None, max_v=None):
"""
原地修改张量 ``x`` ，限制其数值在特定的范围 ``[min_v, max_v]`` 内。对于张量 ``x`` 中的任何一个数值均满足：
.. math::
x_i = \\min(\\max(x_i, \\text{min_value}_i), \\text{max_value}_i)
参数：
- x (Var): Var类型的张量
- min_v (int, optional): 数值的最小允许值。如果为 None，则不应用最小值限制
- max_v (int, optional): 数值的最大允许值。如果为 None，则不应用最大值限制
返回值：
就地修改输入张量 ``x`` ，其数值的维度被限制在 ``[min_v, max_v]`` 之间 。
代码示例：
>>> x = jt.randn(4)
jt.Var([ 0.77959126 0.8394206 -1.426833 0.7886605 ], dtype=float32)
>>> jt.clamp_(x, -0.5, 0.5)
>>> print(x)
jt.Var([ 0.5 0.5 -0.5 0.5], dtype=float32)
注意事项:
- 如果 ``min_v`` 大于 ``max_v`` ，会引起异常
"""
return x.assign(x.clamp(min_v=min_v, max_v=max_v))
Var.clamp_ = clamp_
[文档]
def outer(x, y):
"""
计算两个一维向量的外积。具体公式如下:
.. math::
res_{i,j} = x_{i} \\times y_{j}
参数：
- x(Var, numpy array, python sequence)：输入的 Var 或者 numpy 数组 或者 python 序列，表示第一个向量。
- y(Var, numpy array, python sequence): 输入的 Var 或者 numpy 数组 或者 python 序列，表示第二个向量。
返回值：
两个一维向量的外积(Var)
代码示例:
>>> x = jt.arange(3)
>>> y = jt.arange(4)
>>> jt.outer(x, y)
jt.Var([[0 0 0 0]
[0 1 2 3]
[0 2 4 6]], dtype=int32)
>>> x.outer(y)
jt.Var([[0 0 0 0]
[0 1 2 3]
[0 2 4 6]], dtype=int32)
"""
return jt.multiply(x.unsqueeze(1), y.unsqueeze(0))
Var.outer = outer
[文档]
def erfinv_(x):
"""
函数实现的是数学中的误差反函数(erf)的就地版本 :math:`erf^{-1}(x)`
其中，:math:`erf` 函数定义如下，``x`` 为实数：
.. math::
\\text{erf}(x) = \\frac{2}{\\sqrt{\\pi}} \\int_{0}^{x} e^{-t^{2}}\\, dt
:math:`erfinv` 函数则是 :math:`erf` 函数的反函数.
参数:
x (float): 输入的 Jittor 变量。
返回:
与输入类型和形状相同的 Jittor 变量(Var)。
代码示例:
>>> import jittor as jt
>>> x = jt.array([0.5, 0.0, -0.5])
>>> y = jt.erfinv_(x)
jt.Var([ 0.47693628 0. -0.47693628], dtype=float32)
"""
return x.assign(x.erfinv())
Var.erfinv_ = erfinv_
[文档]
def erf_(x):
"""
函数实现的是 对 ``x`` 实施原址元素级误差函数运算，并返回结果。
误差函数是高斯函数积分，可以表示为：
.. math::
\\text{erf}(x) = \\frac{2}{\\sqrt{\\pi}} \\int_{0}^{x} e^{-t^2} dt
参数:
x(Var)：输入的需要进行原址元素级误差函数运算的 Jittor 变量，可以是标量，向量，矩阵或任意维度的张量。
代码示例：
>>> import jittor as jt
>>> x = jt.array([0, 1., 2.])
>>> print(jt.erf_(x))
jt.Var([0. 0.8427008 0.9953223], dtype=float32)
返回值：
原位计算后的与 ``x`` 具有相同形状和数据类型的 Jittor 变量(Var)。
"""
return x.assign(x.erf())
Var.erf_ = erf_
[文档]
def abs_(x):
"""
计算给定输入的张量 ``x`` 的绝对值，并直接在原地进行修改。数学公式如下:
.. math::
y = |x|
其中 :math:`y` 是输出张量，:math:`x` 是输入张量。
参数:
- x (Var): 要进行绝对值计算的输入张量。
返回值:
计算后的原地修改的张量(Var)，维度和输入相同。
代码示例：
>>> import jittor as jt
>>> x = jt.array([-1, -2, 3])
>>> jt.abs_(x)
jt.Var([1 2 3], dtype=int32)
"""
return x.assign(x.abs())
Var.abs_ = abs_
[文档]
def sigmoid_(x):
"""
原地版本的 sigmoid 函数。这个函数会把输入的 Jittor 变量直接进行替换。
sigmoid 函数的数学形式为:
.. math::
\\sigma(x) = \\frac{1}{1 + e^{-x}}
参数:
x (Var): 输入的 Jittor 变量。
返回值:
应用了 :math:`sigmoid` 函数后的 Jittor 变量(Var)。
代码示例:
>>> import jittor as jt
>>> a = jt.array([1.0, 2.0, 3.0])
>>> b = sigmoid_(a)
jt.Var([0.7310586 0.880797 0.95257413], dtype=float32)
"""
return x.assign(x.sigmoid())
Var.sigmoid_ = sigmoid_
[文档]
def sqrt_(x):
"""
对给定的输入变量 ``x`` 进行开原地开平方，并将结果赋值给原变量。使用算法
.. math::
\\sqrt{x}
对输入的每一个元素进行开平方。
参数:
x (Var): 需要进行开平方的变量。
返回值:
经过开平方处理的变量(Var)。
代码示例:
>>> import jittor as jt
>>> a = jt.array([4., 9.])
>>> jt.sqrt_(a)
jt.Var([2., 3.], dtype=float32)
"""
return x.assign(x.sqrt())
Var.sqrt_ = sqrt_
[文档]
def add_(x, y):
"""
函数是 :math:`add()` 的 原地替换版本,即结果会在原始变量x上就地进行修改，不会产生新的变量。
参数：
- x(Jittor.Var)：第一个输入变量，其值会被函数的结果替代。
- y(Jittor.Var)：第二个输入变量。
返回值：
返回的是经过添加操作后的x(Var)。
代码示例：
>>> import jittor as jt
>>> x = jt.array([1.0, 2.0, 3.0])
>>> y = jt.array([4.0, 5.0, 6.0])
>>> jt.add_(x, y)
jt.Var([5.0, 7.0, 9.0], dtype=float32)
"""
return x.assign(x.add(y))
Var.add_ = add_
[文档]
def multiply_(x, y):
"""
函数是 :math:`multiply()` 的就地版本，即在原始变量上进行操作，不会创建新的变量。该函数主要进行两个输入参数的乘法运算。
数学公式：
:math:`z = x \\times y`
其中 ``x`` 和 ``y`` 都是输入变量。
参数：
- x (int,flot or Jittor.Variable): 输入变量。
- y (int,flot or Jittor.Variable)：输入变量。
代码示例：
>>> import jittor as jt
>>> x = jt.array([1.0, 2.0, 3.0])
>>> y = jt.array([4.0, 5.0, 6.0])
>>> z = jt.multiply_(x, y)
>>> print(z)
array([4., 10., 18.], dtype=float32)
返回值：
output(Var)： ``x`` 和 ``y`` 的乘法结果。它的形状和数据类型与输入变量 ``x`` 相同。
"""
return x.assign(x.multiply(y))
Var.multiply_ = multiply_
[文档]
def type_as(a, b):
"""
将张量 ``a`` 的数据类型修改为与张量 ``b`` 的数据类型相同。
参数:
- a (jt.Var) : 输入的源张量，需要修改其数据类型的张量。
- b (jt.Var) : 输入的目标张量，将会根据此张量的数据类型修改a的数据类型。
返回值:
``a`` 中数据类型已经修改为张量 ``b`` 的数据类型的新张量。类型为jittor.Var。
代码示例：
>>> a = jt.float32([1,2,3])
>>> b = jt.int32([4,5,6])
>>> a = jt.type_as(a, b)
>>> print(a)
jt.Var([1 2 3], dtype=int32)
"""
return a.unary(op=b.dtype)
Var.type_as = type_as
# Var.astype = Var.cast
[文档]
def masked_fill(x, mask, value):
"""
对张量 ``x`` 中 ``mask`` 为 ``True`` 的位置进行填充值 ``value``。需要保证 ``mask`` 的形状必须和 ``x`` 相匹配。
参数:
- x (jt.Var) : 输入的张量。
- mask (bool Var) : 和 ``x`` 形状相同的布尔类型张量。
- value (float) : 用于填充张量 ``x`` 中对应 ``mask`` 为 ``True`` 的位置的值。
代码示例：
>>> x = jt.array([1, 2, 3])
>>> mask = jt.array([False, True, False])
>>> print(jt.ops.masked_fill(x, mask, -1))
jt.Var([ 1 -1 3], dtype=int32)
返回值:
``x`` 填充后的张量。类型为jittor.Var。
"""
return jt.ternary(mask, value, x)
Var.masked_fill = masked_fill
[文档]
def sqr(x):
"""
创建一个张量，其将 ``x`` 中的每一个数值进行平方运算。
.. math::
out_i = \sqr{input_i}
参数：
- x (Var): Var类型的张量
返回值：
返回一个新的张量, 其数值是 ``x`` 中对应位置进行平方运算后的值
代码示例：
>>> x = jt.randn(5)
>>> x
jt.Var([-0.29049814 0.3832289 -0.6810643 -0.65432096 0.39759803], dtype=float32)
>>> jt.sqr(x)
jt.Var([0.08438917 0.14686438 0.4638486 0.4281359 0.1580842 ], dtype=float32)
"""
return x*x
Var.sqr = sqr
[文档]
def pow(x, y):
'''
对两个张量逐个元素计算乘方，也可以使用 ``**`` 运算符调用
参数：
- x (Var): 乘方运算的底数
- y (Var): 乘方运算的幂次，形状与 ``x`` 相同
返回值：
乘方结果(Var)，形状与 ``x`` 和 ``y`` 相同
代码示例：
>>> x = jt.array([1, 2, 3])
>>> y = jt.array([2, 3, 4])
>>> jt.pow(x, y)
jt.Var([ 1 8 81], dtype=int32)
>>> x ** y
jt.Var([ 1 8 81], dtype=int32)
'''
if isinstance(x,Var) and isinstance(y, (ori_int, ori_float)) and y == 2:
return x.sqr()
return core.ops.pow(x, y)
Var.pow = Var.__pow__ = pow
[文档]
def argmax(x: Var, dim: int, keepdims:bool=False):
"""
用于找出张量 ``x`` 中某个维度上元素的最大值的索引。
参数：
- x (Var): Var类型的张量
- dim (int， optional): 指定在哪个维度上进行求最大值, 如果 ``dim`` 指定为 ``None``, 则张量自动被展开
- keepdims (bool, optional): 是否保持原Var的维度，默认为 ``False``
返回值：
返回一个新的张量，其中包含指定维度上每个“切片”中最大元素的索引。
代码示例：
>>> x = jt.randn(3,2)
jt.Var([[-0.1429974 -1.1169171 ]
[-0.35682714 -1.5031573 ]
[ 0.66668254 1.1606413 ]], dtype=float32)
>>> jt.argmax(x, 0)
(jt.Var([2 2], dtype=int32), jt.Var([0.66668254 1.1606413 ], dtype=float32))
注意事项：
- 如果在指定的维度上有多个最大值，``argmax`` 函数通常会返回这些最大值中第一个遇到的索引
"""
if dim is None:
dim = 0
x = x.flatten()
return jt.arg_reduce(x, "max", dim, keepdims)
Var.argmax = argmax
[文档]
def argmin(x, dim: int, keepdims:bool=False):
"""
返回输入张量中沿给定维度的最小值的索引。
参数:
- x (jt.Var) : 输入的张量。
- dim(int) : 计算最小值的索引的维度。
- keepdims (bool) : 是否要求输出张量的维度是否与输入张量保持一致。如果这个值为True,结果张量的维度将与输入张量 ``x`` 的维度一致。默认值: ``False`` 。
返回值:
包含两个张量(Var)的元组(tuple)，第一位为返回输入张量 ``x`` 沿给定维度 ``dim`` 的最小值的索引的张量，第二位为索引对应的最小值的张量。类型为Var的元组。
代码示例：
>>> x = jt.random((4,4))
jt.Var([[0.4455578 0.4899658 0.854926 0.7409189 ]
[0.62386227 0.2531969 0.48017916 0.37137902]
[0.76723385 0.8991033 0.22901663 0.08265838]
[0.23934306 0.63881433 0.55240405 0.2548534 ]], dtype=float32)
>>> jt.argmin(x, dim=0)
(jt.Var([3 1 2 2], dtype=int32), jt.Var([0.23934306 0.2531969 0.22901663 0.08265838], dtype=float32))
>>> jt.argmin(x, dim = 1)
(jt.Var([0 1 3 0], dtype=int32), jt.Var([0.4455578 0.2531969 0.08265838 0.23934306], dtype=float32))
>>> jt.argmin(x, dim = 1, keepdims = True)
(jt.Var([[0]
[1]
[3]
[0]], dtype=int32),
jt.Var([[0.4455578 ]
[0.2531969 ]
[0.08265838]
[0.23934306]], dtype=float32))
"""
return jt.arg_reduce(x, "min", dim, keepdims)
Var.argmin = argmin
[文档]
def randn(*size, dtype="float32", requires_grad=True) -> Var:
"""
返回一个具有标准正态分布的张量。
参数:
- size (int or a sequence of int) : 生成张量的形状。
- dtype(str,optional) : 生成张量的数据类型，默认值: ``float32`` 。
- requires_grad (bool) : 是否能够进行梯度反向传播。默认值: ``True`` 。
返回值:
生成的具有标准正态分布的张量。类型为jittor.Var。
代码示例：
>>> jt.randn(3)
jt.Var([-1.019889 -0.30377278 -1.4948598 ], dtype=float32)
>>> jt.randn(2, 3)
jt.Var([[-0.15989183 -1.5010914 0.5476955 ]
[-0.612632 -1.1471151 -1.1879086 ]], dtype=float32)
"""
if isinstance(size, tuple) and isinstance(size[0], (tuple, list, NanoVector)): size = size[0]
arr = jt.random(size, dtype, "normal")
if not requires_grad: return arr.stop_grad()
return arr
[文档]
def rand(*size, dtype="float32", requires_grad=True) -> Var:
"""
返回一个元素值用从[0, 1)均匀分布中抽样取得的随机Var形式张量。
参数:
- \\*size (序列的整数) -- 输出张量的期望形状，例如 (m, n, k)。
- dtype (str, 可选) -- 输出张量的数据类型，默认为 ``float32`` 。
- requires_grad (bool, 可选) -- 是否需要返回梯度，默认为 ``False`` 。
返回值:
返回一个 \\*size 维Var形式张量，元素值用从[0, 1)均匀分布中抽样取得的随机Var形式张量。
代码示例：
>>> import jittor as jt
>>> jt.random.seed(3)
>>> jt.rand(3,3)
jt.Var([[0.5509497 , 0.6197729 , 0.909781 ],
[0.04279451, 0.10358319, 0.9024445 ],
[0.09957159, 0.95744824, 0.546594 ]], dtype=float32)
"""
if isinstance(size, tuple) and isinstance(size[0], (tuple, list, NanoVector)): size = size[0]
arr = jt.random(size, dtype)
if not requires_grad: return arr.stop_grad()
return arr
[文档]
def rand_like(x, dtype=None) -> Var:
"""
生成一个元素值从[0,1]均匀分布中抽样取得的张量，并且形状与输入的张量形状相同。
参数:
- x (jt.Var) : 输入的张量。
- dtype (str,optional) : 生成的随机张量的数据类型。如果此参数为None, 则生成的张量的数据类型将与输入的张量的数据类型相同。默认值: ``None`` 。
返回值:
与输入的张量 ``x`` 形状相同，值为[0,1]均匀分布随机抽样的张量。
代码示例：
>>> x = jt.zeros((2, 3))
>>> jt.rand_like(x)
jt.Var([[0.6164821 0.21476883 0.61959815]
[0.58626485 0.35345772 0.5638483 ]], dtype=float32)
"""
if dtype is None: dtype = x.dtype
return jt.random(x.shape, dtype)
[文档]
def randn_like(x, dtype=None) -> Var:
"""
生成一个元素服从标准正态分布的张量，并且形状与输入的张量形状相同的随机张量。
参数:
- x (jt.Var) : 输入的张量。
- dtype (str,optional) : 生成的随机张量的数据类型。如果此参数为None, 则生成的张量的数据类型将与输入的张量的数据类型相同。默认值: ``None`` 。
返回值:
与输入的张量 ``x`` 形状相同，元素服从标准正态分布的张量(Var)。
代码示例：
>>> x = jt.zeros((2, 3))
>>> jt.randn_like(x)
jt.Var([[-1.1647032 0.34847224 -1.3061888 ]
[ 1.068085 -0.34366122 0.13172573]], dtype=float32)
"""
if dtype is None: dtype = x.dtype
return jt.random(x.shape, x.dtype, "normal")
[文档]
def randint(low, high=None, shape=(1,), dtype="int32") -> Var:
"""
用于在一个指定的范围 [ ``low`` , ``high``)内生成随机整数张量。
参数:
- low (int, optional) : 生成随机数的范围的下限。默认值: ``0`` 。
- high (int) : 生成随机数的范围的上限。
- shape (tuple, optional) : 生成的随机张量的形状。默认值: ``(1,)`` 。
- dtype (str,optional) : 生成的随机张量的数据类型。默认值: ``int32`` 。
返回值:
生成的随机张量(Var)。
代码示例：
>>> jt.randint(3, shape=(3, 3))
jt.Var([[2 0 2]
[2 1 2]
[2 0 1]], dtype=int32)
>>> jt.randint(1, 3, shape=(3, 3))
jt.Var([[2 2 2]
[1 1 2]
[1 1 1]], dtype=int32)
"""
if high is None: low, high = 0, low
v = (jt.random(shape) * (high - low) + low).clamp(low, high-0.5)
v = jt.floor_int(v)
return v.astype(dtype)
[文档]
def randint_like(x, low, high=None) -> Var:
"""
用于在一个指定的范围[ ``low`` , ``high``)内生成随机整数张量，并且形状与输入变量 ``x`` 相同。
参数:
- x (jt.Var) : 输入的张量。
- low (int, optional) : 生成随机数的范围的下限。默认值: ``0`` 。
- high (int) : 生成随机数的范围的上限。
返回值:
生成的形状与输入变量 ``x`` 相同的随机张量(Var)。
代码示例：
>>> x = jt.zeros((2, 3))
>>> jt.randint_like(x, 10)
jt.Var([[9. 3. 4.]
[4. 8. 5.]], dtype=float32)
>>> jt.randint_like(x, 10, 20)
jt.Var([[17. 11. 18.]
[14. 17. 15.]], dtype=float32)
"""
return randint(low, high, x.shape, x.dtype)
[文档]
def normal(mean, std, size=None, dtype="float32") -> Var:
"""
生成产生满足正态分布的随机数的张量。
参数:
- mean (int or jt.Var) : 正态分布的均值，是标量或者一个与需要生成的数组shape相同的张量。
- std (int or jt.Var) : 正态分布的标准差，是标量或者一个与需要生成的数组shape相同的张量。
- size (tuple, optional) : 生成的随机张量的形状。默认值: ``None`` 。
- dtype (str,optional) : 生成的随机张量的数据类型。默认值: ``float32`` 。
返回值:
生成的产生满足正态分布的随机数的张量(Var)。
代码示例：
>>> jt.normal(5, 3, size=(2,3))
jt.Var([[ 8.070848 7.654219 10.252696 ]
[ 6.383718 7.8817277 3.0786133]], dtype=float32)
>>> mean = jt.randint(low=0, high=10, shape=(10,))
>>> jt.normal(mean, 0.1)
jt.Var([1.9524184 1.0749301 7.9864206 5.9407325 8.1596155 4.824019 7.955083
8.972998 6.0674286 8.88026 ], dtype=float32)
"""
if size is None:
if isinstance(mean, Var) and isinstance(std, Var):
assert mean.shape == std.shape
size = mean.shape
else:
if isinstance(mean, Var): size = mean.shape
if isinstance(std, Var): size = std.shape
return jt.init.gauss(size, dtype, mean, std)
[文档]
def attrs(var):
"""
获取变量 ``var`` 的属性。
参数:
- var (jt.Var) : 输入的变量。
返回值:
返回一个字典，包含变量的所有属性。
代码示例：
>>> a = jt.array([1, 2, 3])
>>> print(jt.attrs(a))
{'is_stop_fuse': False, 'is_stop_grad': True, 'shape': [3,], 'dtype': int32}
"""
return {
"is_stop_fuse": var.is_stop_fuse(),
"is_stop_grad": var.is_stop_grad(),
"shape": var.shape,
"dtype": var.dtype,
}
Var.attrs = attrs
[文档]
def fetch(*args):
'''
注册一个在每个张量都计算完成并拿到数据后的回调函数。给出若干个参数，里面是 ``Var`` 的都计算完成后，调用回调函数。传入回调函数的参数与 ``fetch`` 除去最后回调函数的参数对应，但是所有 ``Var`` 都替换为其 ``data`` 。
参数：
- 前若干个参数为等待完成的 ``Var`` 和其它参数，最后一个参数是回调函数
返回值：
无
代码示例：
>>> a = jt.zeros((2, 3), dtype='float32')
>>> jt.fetch(a, lambda x: print('got', repr(x)))
>>> jt.sync_all()
got array([[0., 0., 0.],
[0., 0., 0.]], dtype=float32)
'''
assert len(args)>=1
func = args[-1]
assert callable(func)
args = list(args[:-1])
if len(args)>0 and isinstance(args[0], Sequence) \
and len(args[0])>=1 and isinstance(args[0][0], Var):
raise TypeError("jt.Var should not inside a list or tuple.")
var_map = []
variables = []
for i, v in enumerate(args):
if isinstance(v, Var):
variables.append(v)
var_map.append(i)
args[i] = None
def callback(*results):
for i,v in enumerate(results):
args[var_map[i]] = v
func(*args)
core.ops.fetch(variables, callback)
Var.fetch = fetch
[文档]
def display_memory_info():
"""
展示当前Jittor正在使用的内存相关信息。
参数:
不需要任何参数
代码示例：
>>> jt.display_memory_info()
=== display_memory_info ===
total_cpu_ram: 125.7GB total_device_ram: 23.69GB
hold_vars: 0 lived_vars: 0 lived_ops: 0
name: sfrl is_device: 1 used: 0 B(-nan%) unused: 0 B(-nan%) total: 0 B
name: sfrl is_device: 0 used: 0 B(-nan%) unused: 0 B(-nan%) total: 0 B
name: sfrl is_device: 0 used: 0 B(-nan%) unused: 0 B(-nan%) total: 0 B
name: sfrl is_device: 0 used: 0 B(-nan%) unused: 0 B(-nan%) total: 0 B
cpu&gpu: 0 B gpu: 0 B cpu: 0 B
free: cpu(51.98GB) gpu(22.28GB)
swap: total( 0 B) last( 0 B)
===========================
返回值:
None。
"""
import inspect, os
f = inspect.currentframe()
fileline = inspect.getframeinfo(f.f_back)
fileline = f"{os.path.basename(fileline.filename)}:{fileline.lineno}"
core.display_memory_info(fileline)
[文档]
def load(path: str):
"""
从指定路径 ``path`` 加载一个模型。
参数:
- path (str) :需要加载模型的具体路径，此参数为必需。
返回值:
返回加载后的模型。
代码示例：
>>> model = jt.load(\"/path/to/model\")
"""
model_dict = safeunpickle(path)
return model_dict
[文档]
def save(params_dict, path: str):
"""
将指定的参数 ``params_dict`` 保存到指定的路径 ``path`` 。
参数:
- params_dict (list or dictionary) : 待保存的参数。
- path (str) : 保存的文件路径。
返回值:
None。
代码示例：
>>> params_dict = {\"weight\": np.array([1, 2, 3]), \"bias\": np.array([0.1, 0.2, 0.3])}
>>> jt.save(params_dict, \"params.p\")
"""
safepickle(params_dict, path)
def _uniq(x):
"""
该函数用于去除列表中的重复元素。如果输入列表的元素为可变对象（如列表,字典等), 则使用object的id进行去重。
参数：
x(list): 输入的列表，其中的元素可能重复。
代码示例：
>>> inputlist=[1,1,2,2,3,3]
>>>_uniq(inputlist)
[1,2,3]
返回值：
去除重复元素后的列表(list)，顺序保持与原列表中的第一次出现顺序一致。
"""
a = set()
b = []
for i in x:
j = id(i)
if j not in a:
a.add(j)
b.append(i)
return b
[文档]
class Module:
"""
Module 类通常用于构建神经网络模型。Module 是 Jittor 深度学习框架的基础模块，是所有神经网络层和模型的基类。提供了构建神经网络模型所需的各种基础方法，包括前向计算( execute 方法)、设置混合精度计算(比如 float16/32 方法)、设置模型的训练/评估状态( train/eval 方法)、获取模型参数( parameters 方法) 等等。同时，由于 Jittor 是一个 JIT 编译框架，它还使用了很多 C++ 代码和编译技术来提高执行效率，这使得 Jittor 的 Module 类在运行速度上具有优势。
参数：
具体的参数取决于所定义的模型结构和网络架构。主要包括以下方面：
- 初始化参数：初始化 Module 对象的参数。例如权重初始化方式、偏置项初始化方式等。
- 超参数：定义模型结构和行为的参数。例如隐藏层的维度、激活函数的类型、优化器的学习率等。
- 输入数据相关参数：定义输入数据的特征维度、输入数据的形状等。
- 其他配置参数：指定其他模型相关的配置。例如是否使用 GPU 加速、是否打印训练过程中的日志等。
形状：
- Input：输入形状根据具体的模型结构和网络架构而定。
- Output：输出形状根据具体的模型结构和网络架构而确定。
代码示例：
>>> import jittor as jt
>>> from jittor import nn
>>> class LinearModel(nn.Module):
... def __init__(self):
... self.linear = nn.Linear(10, 2)
... def execute(self, x):
... x = self.linear(x)
... return x
>>> net = LinearModel()
>>> x = jt.random((10, 10))
>>> out = net(x)
>>> print(out.shape)
[10,2,]
"""
def __init__(self, *args, **kw):
pass
def execute(self, *args, **kw):
''' Executes the module computation.
Raises NotImplementedError if the subclass does not override the method.
'''
raise NotImplementedError("Please implement 'execute' method of "+str(type(self)))
def __call__(self, *args, **kw):
return self.execute(*args, **kw)
def __repr__(self):
return self.__str__()
def _get_name(self):
return self.__class__.__name__
def __name__(self):
pass
def dfs(self, parents, k, callback, callback_leave=None, recurse=True):
''' 遍历模块的实用函数。 '''
n_children = 0
for v in self.__dict__.values():
if isinstance(v, Module):
n_children += 1
ret = callback(parents, k, self, n_children)
if ret == False: return
if recurse:
for k,v in self.__dict__.items():
if not isinstance(v, Module):
continue
parents.append(self)
v.dfs(parents, k, callback, callback_leave)
parents.pop()
if callback_leave:
callback_leave(parents, k, self, n_children)
def __str__(self):
ss = []
def callback(parents, k, v, n):
# indent key:class_name(extra_repr)
k = f"{k}: " if k is not None else ""
s = f"{' '*(len(parents)*4)}{k}{v.__class__.__name__}"
if n:
s += '('
else:
s += f"({v.extra_repr()})"
ss.append(s)
def callback_leave(parents, k, v, n):
if n:
ss.append(' '*(len(parents)*4)+')')
self.dfs([], None, callback, callback_leave)
return "\n".join(ss)
[文档]
def parameters(self, recurse=True) -> List:
''' 返回模块参数的列表。
示例代码::
>>> net = nn.Sequential(nn.Linear(2, 10), nn.ReLU(), nn.Linear(10, 2))
>>> for p in net.parameters():
... print(p.name)
...
>>> for p in net.parameters():
... print(p.name())
...
0.weight
0.bias
2.weight
2.bias
'''
ps = []
stack = []
def callback(parents, k, v, n):
stack.append(str(k))
dc = v.__dict__
if isinstance(v, nn.ParameterList):
dc = v.params
for k2, p in dc.items():
if isinstance(k2, str) and k2.startswith("_"): continue
if isinstance(p, Var):
ps.append(p)
pname = ".".join(stack[1:]+[str(k2)])
if len(pname) > len(p.name()):
p.name(pname)
def callback_leave(parents, k, v, n):
stack.pop()
self.dfs([], None, callback, callback_leave, recurse)
return _uniq(ps)
def state_dict(self, to=None, recurse=True):
''' 返回一个包含模块及其后代的Jittor Var的字典。
Args:
to: target type of var, canbe None or 'numpy' or 'torch'
Return:
dictionary of module's states.
Example::
import jittor as jt
from jittor.models import resnet50
jittor_model = resnet50()
dict = jittor_model.state_dict()
jittor_model.load_state_dict(dict)
Example2(export Jittor params to PyTorch)::
import jittor as jt
from jittor.models import resnet50
jittor_model = resnet50()
import torch
from torchvision.models import resnet50
torch_model = resnet50()
torch_model.load_state_dict(jittor_model.state_dict(to="torch"))
'''
uniq_set = set()
ps = {}
stack = []
def callback(parents, k, v, n):
stack.append(str(k))
dc = v.__dict__
if isinstance(v, nn.ParameterList):
dc = v.params
for k2, p in dc.items():
if isinstance(k2, str) and k2.startswith("_"): continue
if isinstance(p, Var):
if id(p) in uniq_set: continue
if not getattr(p, "persistent", True):
continue
uniq_set.add(id(p))
pname = ".".join(stack[1:]+[str(k2)])
ps[pname] = p
if len(pname) > len(p.name()):
p.name(pname)
def callback_leave(parents, k, v, n):
stack.pop()
self.dfs([], None, callback, callback_leave, recurse)
if to == "numpy":
for k,v in ps.items():
if isinstance(v, Var):
ps[k] = v.numpy()
elif to == "torch":
import torch
for k,v in ps.items():
if isinstance(v, Var):
ps[k] = torch.Tensor(v.numpy())
return ps
[文档]
def named_parameters(self, recurse=True) -> List[Tuple[str, Var]]:
''' 返回模块参数及其名称的列表。
示例代码::
>>> net = nn.Linear(2, 5)
>>> net.named_parameters()
[('weight', jt.Var([[ 0.5964666 -0.3175258 ]
[ 0.41493994 -0.66982657]
[-0.32677156 0.49614117]
[-0.24102807 -0.08656466]
[ 0.15868133 -0.12468725]], dtype=float32)),
('bias', jt.Var([-0.38282675 0.36271113 -0.7063226 0.02899247 0.52210844], dtype=float32))]
'''
state_dict = self.state_dict(recurse=recurse)
return list(state_dict.items())
[文档]
def load_state_dict(self, params) -> None:
'''
从字典中加载模块的参数。
'''
self.load_parameters(params)
def _load_from_state_dict(self, state, prefix="", *args, **kw):
if len(prefix):
new_state = {}
for k,v in state.items():
if k.startswith(prefix):
new_state[k[len(prefix):]] = v
state = new_state
self.load_state_dict(state)
[文档]
def cuda(self, device=None):
'''
将模块放置于cuda运行
'''
flags.use_cuda = 1
return self
def npu(self, device=None):
flags.use_cuda = 1
return self
[文档]
def modules(self) -> List:
''' 递归地返回模块中子模块的列表。
示例代码::
>>> net = nn.Sequential(nn.Linear(2, 10), nn.ReLU(), nn.Linear(10, 2))
>>> net.modules()
[Sequential(
0: Linear(2, 10, float32[10,], None)
1: relu()
2: Linear(10, 2, float32[2,], None)
), Linear(2, 10, float32[10,], None), relu(), Linear(10, 2, float32[2,], None)]
'''
ms = []
def callback(parents, k, v, n):
if isinstance(v, Module):
ms.append(v)
self.dfs([], None, callback, None)
return _uniq(ms)
[文档]
def named_modules(self):
''' 递归地返回子模块及其名称的列表。
代码示例::
>>> net = nn.Sequential(nn.Linear(2, 10), nn.ReLU(), nn.Linear(10, 2))
>>> net.named_modules()
[('', Sequential(
0: Linear(2, 10, float32[10,], None)
1: relu()
2: Linear(10, 2, float32[2,], None)
)), ('0', Linear(2, 10, float32[10,], None)), ('1', relu()), ('2', Linear(10, 2, float32[2,], None))]
'''
ms = []
stack = []
def callback(parents, k, v, n):
if isinstance(v, Module):
stack.append(str(k))
name = ".".join(stack[1:])
ms.append((name, v))
def callback_leave(parents, k, v, n):
stack.pop()
self.dfs([], "", callback, callback_leave)
return ms
def add_module(self, name, module):
setattr(self, name ,module)
return self
@property
def _modules(self):
return { k:v for k,v in self.__dict__.items() if isinstance(v, Module) }
@property
def _parameters(self):
return { k:v for k,v in self.__dict__.items() if isinstance(v, Var) }
[文档]
def requires_grad_(self, requires_grad=True):
'''为所有参数和子模块设置 requires_grad。 '''
self._requires_grad = requires_grad
self._place_hooker()
return self
def __hooked_call__(self, *args, **kw):
if hasattr(self, "__fhook2__"):
if len(kw):
self.__fhook2__(self, args, kw)
else:
self.__fhook2__(self, args)
if hasattr(self, "__bihook__"):
if len(kw):
LOG.w("backward hook not support kw")
args = grad_hooker(args, self.__bihook__)
if hasattr(self, "_requires_grad") and not self._requires_grad:
with jt.no_grad():
ret = self.__hooked_call__(*args, **kw)
else:
ret = self.__hooked_call__(*args, **kw)
if hasattr(self, "__bohook__"):
if len(kw):
LOG.w("backward hook not support kw")
if isinstance(ret, Var):
ret = grad_hooker((ret,), self.__bohook__)[0]
else:
ret = grad_hooker(ret, self.__bohook__)
if hasattr(self, "__fhook__"):
if len(kw):
self.__fhook__(self, args, ret, kw)
else:
self.__fhook__(self, args, ret)
return ret
def _place_hooker(self):
cls = self.__class__
if hasattr(cls, "__hooked__"):
return
cls.__hooked__ = True
cls.__call__, cls.__hooked_call__ = \
cls.__hooked_call__, cls.__call__
[文档]
def register_forward_hook(self, func):
''' 注册一个前向函数钩子，在Module.execute之后调用。
此函数可以按照如下格式调用::
hook(module, input_args, output)
或者::
hook(module, input_args, output, input_kwargs)
'''
self.__fhook__ = func
self._place_hooker()
[文档]
def remove_forward_hook(self):
''' 移除当前的前向钩子。 '''
if hasattr(self,"__fhook__"):
delattr(self,"__fhook__")
[文档]
def register_pre_forward_hook(self, func):
''' 注册一个前向函数钩子，在Module.execute之前调用。
此函数可以按照如下格式调用::
hook(module, input_args)
或者::
hook(module, input_args, input_kwargs)
'''
self.__fhook2__ = func
self._place_hooker()
[文档]
def remove_pre_forward_hook(self):
''' 移除当前的预前向钩子。 '''
if hasattr(self,"__fhook2__"):
delattr(self,"__fhook2__")
def register_input_backward_hook(self, func):
self.__bihook__ = func
self._place_hooker()
def remove_input_backward_hook(self):
if hasattr(self,"__bihook__"):
delattr(self,"__bihook__")
def register_output_backward_hook(self, func):
self.__bohook__ = func
self._place_hooker()
def remove_output_backward_hook(self):
if hasattr(self,"__bohook__"):
delattr(self,"__bohook__")
[文档]
def register_backward_hook(self, func):
''' 在此模块的反向传播过程中挂钩输入和输出。
这个函数参数如下所示::
- hook(module, grad_input:tuple(jt.Var)
- grad_output:tuple(jt.Var)) -> tuple(jt.Var) or None
``grad_input`` 是本模块原本的梯度输入。``grad_output`` 是梯度输出, 返回值会替代输入的输入梯度。
'''
_grad_output = None
def bohook(grad_output):
nonlocal _grad_output
_grad_output = grad_output
def bihook(grad_input):
return func(self, grad_input, _grad_output)
self.register_input_backward_hook(bihook)
self.register_output_backward_hook(bohook)
[文档]
def remove_backward_hook(self):
''' 删除反向传播的输入和输出hook
'''
self.remove_input_backward_hook()
self.remove_output_backward_hook()
[文档]
def children(self) -> List:
''' 返回子模块的列表。 '''
cd = []
def callback(parents, k, v, n):
if len(parents) == 1 and isinstance(v, Module):
cd.append(v)
return False
self.dfs([], None, callback, None)
return cd
def extra_repr(self):
ss = []
n = len(self.__init__.__code__.co_varnames)
if self.__init__.__defaults__ is not None:
n -= len(self.__init__.__defaults__)
for i, k in enumerate(self.__init__.__code__.co_varnames[1:]):
v = getattr(self, k) if hasattr(self, k) else None
if isinstance(v, Var): v = v.peek()
s = f"{k}={v}" if i >= n else str(v)
ss.append(s)
return ", ".join(ss)
def apply(self, func):
''' 递归地对所有子模块应用一个函数。 '''
for m in self.modules():
func(m)
[文档]
def load_parameters(self, params):
''' 将参数加载到模块中。
:param params: 一个记录模型函数类名称和对应参数的字典。
'''
n_failed = 0
for key in params.keys():
v = self
key_ = key.split('.')
end = 0
for k in key_:
if isinstance(v, nn.Sequential):
if (k in v.layers):
v = v[k]
elif k.isdigit() and (ori_int(k) in v.layers):
v = v[ori_int(k)]
else:
end=1
break
else:
if hasattr(v, k):
v = getattr(v, k)
assert isinstance(v, (Module, Var)), \
f"expect a jittor Module or Var, but got <{v.__class__.__name__}>, key: {key}"
else:
end = 1
break
if end == 1:
if not key.endswith("num_batches_tracked"):
n_failed += 1
LOG.w(f'load parameter {key} failed ...')
else:
assert isinstance(v, Var), \
f"expect a jittor Var, but got <{v.__class__.__name__}>, key: {key}"
if isinstance(params[key], np.ndarray) or isinstance(params[key], list):
param = array(params[key])
elif isinstance(params[key], Var):
param = params[key]
else:
# assume is pytorch tensor
param = array(params[key].cpu().detach().numpy())
if param.shape == v.shape:
LOG.v(f'load parameter {key} success ...')
v.update(param)
v.sync(False, False)
else:
n_failed += 1
LOG.e(f'load parameter {key} failed: expect the shape of {key} to be {v.shape}, but got {param.shape}')
if n_failed:
LOG.w(f"load total {len(params)} params, {n_failed} failed")
[文档]
def save(self, path: str):
''' 将参数保存到文件中。
:param path: path to save.
:type path: str
示例代码::
>>> class Net(nn.Module):
>>> ...
>>> net = Net()
>>> net.save('net.pkl')
>>> net.load('net.pkl')
'''
params = self.state_dict()
safepickle(params, path)
[文档]
def load(self, path: str):
''' 从文件中加载参数。
:param path: 模型位置
:type path: str
代码示例::
>>> class Net(nn.Module):
>>> ...
>>> net = Net()
>>> net.save('net.pkl')
>>> net.load('net.pkl')
本方法也支持从pytorch .pth文件中加载状态字典。
.. note::
当载入的参数与模型定义不一致时, jittor 会输出错误信息, 但是不会抛出异常.
若载入参数出现模型定义中没有的参数名, 则会输出如下信息, 并忽略此参数:
>>> [w 0205 21:49:39.962762 96 __init__.py:723] load parameter w failed ...
若载入参数的 shape 与模型定义不一致, 则会输出如下信息, 并忽略此参数:
>>> [e 0205 21:49:39.962822 96 __init__.py:739] load parameter w failed: expect the shape of w to be [1000,100,], but got [3,100,100,]
如载入过程中出现错误, jittor 会输出概要信息, 您需要仔细核对错误信息
>>> [w 0205 21:49:39.962906 96 __init__.py:741] load total 100 params, 3 failed
'''
self.load_parameters(load(path))
[文档]
def eval(self):
''' 将模块设置为评估eval模式。 '''
def callback(parents, k, v, n):
if isinstance(v, Module):
v.is_train = False
self.dfs([], None, callback, None)
# backup stop grad or not
if not hasattr(self, "backup_grad_state"):
self.backup_grad_state = {}
for p in self.parameters():
if id(p) not in self.backup_grad_state:
self.backup_grad_state[id(p)] = not p.is_stop_grad()
p.stop_grad()
return self
[文档]
def train(self):
''' 将模块设置为评估training模式。 '''
def callback(parents, k, v, n):
if isinstance(v, Module):
v.is_train = True
self.dfs([], None, callback, None)
# backup stop grad or not
if hasattr(self, "backup_grad_state"):
for p in self.parameters():
if id(p) in self.backup_grad_state and self.backup_grad_state[id(p)]:
p.start_grad()
return self
[文档]
def is_training(self) -> bool:
''' 返回模块是否处于training模式。'''
if not hasattr(self, "is_train"):
self.is_train = True
return self.is_train
@property
def training(self):
if not hasattr(self, "is_train"):
self.is_train = True
return self.is_train
@training.setter
def training(self, value):
self.is_train = value
def mpi_param_broadcast(self, root=0):
if not in_mpi: return
for p in self.parameters():
p.update(p.mpi_broadcast(root))
def __setattr__(self, key, value):
object.__setattr__(self, key, value)
def __getattr__(self, key):
return object.__getattribute__(self, key)
def register_buffer(self, key, value, persistent=True):
value.persistent = persistent
object.__setattr__(self, key, value)
return value
@property
def _buffers(self):
buffers = {}
for k,v in self.__dict__.items():
if isinstance(v, jt.Var):
buffers[k] = v
return buffers
def named_buffers(self,recurse=False):
buffers = []
for k,v in self.__dict__.items():
if isinstance(v, jt.Var):
buffers.append((k,v))
return buffers
def named_children(self,):
childs = []
for k,v in self.__dict__.items():
if isinstance(v,Module):
childs.append((k,v))
return childs
[文档]
def float64(self):
'''将所有参数转换为float16。'''
self._amp_level = 0
for p in self.parameters():
if p.dtype.is_float():
p.assign(p.float64())
return self
[文档]
def float32(self):
'''将所有参数转换为float32。'''
self._amp_level = 0
for p in self.parameters():
if p.dtype.is_float():
p.assign(p.float32())
return self
[文档]
def float16(self):
'''将所有参数转换为float16。'''
# self._amp_level = 3 if flags.th_mode else 4
# amp level better set globally
self._amp_level = -1
if self._amp_level >= 0:
cls = self.__class__
cls.__call__ = cls.__half_call__
for p in self.parameters():
if p.dtype.is_float():
p.assign(p.float16())
return self
[文档]
def bfloat16(self):
'''将所有参数转换为bfloat16。'''
# self._amp_level = 3 if flags.th_mode else 4
# amp level better set globally
self._amp_level = -1
if self._amp_level >= 0:
cls = self.__class__
cls.__call__ = cls.__half_call__
for p in self.parameters():
if p.dtype.is_float():
p.assign(p.bfloat16())
return self
def __half_call__(self, *args, **kw):
amp_level = getattr(self, "_amp_level", -1)
if amp_level >= 0:
with flag_scope(amp_level=amp_level):
return self.execute(*args, **kw)
else:
return self.execute(*args, **kw)
def half(self):
'''convert all parameters to float16'''
return self.float16()
[文档]
def float_auto(self):
'''根据 jt.flags.auto_mixed_precision_level 和 jt.flags.amp_reg 自动将所有参数转换为 float16 或 float32。'''
self._amp_level = -1
for p in self.parameters():
if p.dtype.is_float():
p.assign(p.float_auto())
return self
[文档]
class Function(Module):
"""
用于自定义反向操作（backward operations）的函数模块。``Function`` 类继承自 Jittor 的 ``Module`` 类。
其派生类应实现 ``execute`` 和 ``grad`` 两个方法，作用分别是执行前向计算和计算反向梯度。当使用自定义算子进行前向计算时，应调用类方法 ``apply``，而不是直接调用 ``execute``。注意 ``grad`` 接收的参数和返回值都可为 ``None``，表示无梯度。
代码示例：
>>> import jittor as jt
>>> from jittor import Function
>>> class MyFunc(Function):
... def execute(self, x, y):
... self.x = x
... self.y = y
... return x*y, x/y
... def grad(self, grad0, grad1):
... return grad0 * self.y, grad1 * self.x
...
>>> a = jt.array(3.0)
>>> b = jt.array(4.0)
>>> func = MyFunc.apply
>>> c,d = func(a, b)
>>> da, db = jt.grad(c+d*3, [a, b])
>>> print(da,db)
jt.Var([4.], dtype=float32) jt.Var([9.], dtype=float32)
"""
def __call__(self, *args):
if flags.no_grad:
return self.execute(*args)
backup = args
args = list(args)
taped_inputs = []
taped_outputs = []
input_mask = [-1] * len(args)
for i,v in enumerate(args):
if isinstance(v, Var):
if v.is_stop_grad():
# -2 in input_mask represents it is stop_grad
input_mask[i] = -2
continue
v = v.tape()
input_mask[i] = len(taped_inputs)
args[i] = v
taped_inputs.append(v)
ori_res = self.execute(*args)
if not isinstance(ori_res, Sequence):
res = [ori_res]
else:
res = list(ori_res)
output_mask = [-1] * len(res)
for i,v in enumerate(res):
if isinstance(v, Var):
v = v.tape()
output_mask[i] = len(taped_outputs)
res[i] = v
taped_outputs.append(v)
self.input_mask = input_mask
self.output_mask = output_mask
# tape output and input together so
# backward treat them as one operator
tape_together(taped_inputs, taped_outputs, self._grad)
if isinstance(ori_res, Sequence):
return res
else:
return res[0]
def _grad(self, *args):
new_args = ( (args[i] if i>=0 else None) for i in self.output_mask )
ret = self.grad(*new_args)
if not isinstance(ret, Sequence):
ret = (ret,)
new_ret = []
for i, r in enumerate(ret):
j = self.input_mask[i]
if j<0:
# -2 in input_mask represents it is stop_grad
assert r is None or j==-2, f"{type(self)}'s {i}-th returned grad should be None, "\
"because the input value is not jittor variable."
else:
new_ret.append(r)
return new_ret
def dfs(self, parents, k, callback, callback_leave=None, recurse=True):
pass
@classmethod
def apply(cls, *args, **kw):
func = cls()
return func(*args, **kw)
class GradHooker(Function):
def __init__(self, hook):
self.hook = hook
def execute(self, *args):
return args
def grad(self, *grad_input):
ret = self.hook(grad_input)
if ret: grad_input = ret
return grad_input
def grad_hooker(args, hook):
hooker = GradHooker(hook)
return hooker(*args)
[文档]
def register_hook(v, hook):
"""
将钩子函数注册到指定变量上。钩子函数将在变量的backward操作结束后被调用，主要用于获取变量的梯度信息。
参数：
- v (jittor.Var): 需要注册钩子的张量(jittor.Var)
- hook (Callable(jittor.Var, Union{None, jittor.Var})): 用于处理变量梯度的钩子函数，接收Jittor变量作为参数，返回None或Jittor变量。
返回值：
无返回值
代码示例:
>>> import jittor as jt
>>> x = jt.array([0.0, 0.0])
>>> y = x * [1, 2]
>>> y.register_hook(lambda g: g * 2)
>>> dx = jt.grad(y, x)
>>> print(dx)
jt.Var([2. 4.], dtype=float32)
"""
def _hook(grads):
g = hook(grads[0])
if g is not None:
return (g,)
return None
hooker = GradHooker(_hook)
v.swap(hooker(v)[0])
return v
Var.register_hook = register_hook
[文档]
def make_module(func, exec_n_args=1):
"""
用于根据给定的函数和执行参数的数量，创建一个Jittor模块。
参数:
- func (callable): 输入的函数，用于创建Jittor模块。
- exec_n_args (int): 默认为...执行参数的数量。
返回值:
module (Jittor.Module): 根据给定的函数和执行参数的数量创建的Jittor模块。
代码示例:
>>> import jittor as jt
>>> def add_func(a, b):
>>> return a + b
>>> module = jt.make_module(add_func, 2)
>>> print(module)
<class 'jittor.make_module.<locals>.MakeModule'>
注意:
输入的函数必须是可以调用的，并且执行参数的数量必须大于等于0。
"""
class MakeModule(Module):
def __init__(self, *args, **kw):
self.args = args
self.kw = kw
def execute(self, *args):
return func(*args, *self.args, **self.kw)
def __str__(self):
return f"{func.__name__}({self.extra_repr()})"
def extra_repr(self):
return ",".join(map(str, self.args))
MakeModule.__name__ = func.__name__
return MakeModule
[文档]
def dirty_fix_pytorch_runtime_error():
"""
用于处理PyTorch的运行时错误。
注意：
这个方法没有任何参数，对于调用它的客户端是透明的。仅在需要解决PyTorch的运行时错误时使用。没有给出具体的错误类型和解决方法，因为这是由底层库定义的，并且可能会随着时间的推移而改变。
代码示例:
>>> import jittor as jt
>>> jt.dirty_fix_pytorch_runtime_error()
>>> import torch
警告：
这个方法仅作为临时解决方案，使用它可能会影响PyTorch的正常使用。在无法通过其他方式解决的情况下使用。
"""
import os, platform
if platform.system() == 'Linux':
os.RTLD_GLOBAL = os.RTLD_GLOBAL | os.RTLD_DEEPBIND
import jittor_utils
with jittor_utils.import_scope(os.RTLD_GLOBAL | os.RTLD_NOW):
import torch
import atexit
[文档]
class ExitHooks(object):
"""
该类主要用于捕获程序退出时的退出码和异常信息。在 Jittor 中，该类经常被用于添加一些在程序结束时需要进行的操作，比如资源清理等。其使用方式主要是将需要进行的操作添加到 exit 方法中。Jittor 中定义了该类的全局对象 ``jt.hooks`` 供使用。
属性:
- exit_code：程序退出时的退出码。默认值：``None``
- exception：程序退出时的异常信息。默认值：``None``
代码示例:
>>> import jittor as jt
>>> import sys, atexit
>>> hooks = jt.ExitHooks()
>>> hooks.hook()
>>> def exitfoo():
... if hooks.exception is not None:
... print(\"捕获到异常:\", hooks.exception)
... if hooks.exit_code is not None:
... print(\"程序退出，退出码为\", hooks.exit_code)
...
>>> atexit.register(exitfoo)
>>> print(1/0) # 捕获到异常: division by zero
>>> sys.exit(1) # 程序退出，退出码为 1
"""
def __init__(self):
self.exit_code = None
self.exception = None
def hook(self):
self._orig_exit = sys.exit
sys.exit = self.exit
sys.excepthook = self.exc_handler
def exit(self, code=0):
self.exit_code = code
self._orig_exit(code)
def exc_handler(self, exc_type, exc, *args):
self.exception = exc
traceback.print_exception(exc_type, exc, *args)
hooks = ExitHooks()
hooks.hook()
[文档]
def jittor_exit():
"""
主要作用是关闭Jittor，即在完成Jittor绑定的任务后，安全地退出Jittor运行环境。此函数无需任何输入参数并且没有返回值。
代码示例：
>>> import jittor as jt
>>> jt.jittor_exit()
"""
if hooks.exit_code is not None:
pass
elif hooks.exception is not None:
pass
else:
pass
# core.sync_all(True)
core.cleanup()
atexit.register(jittor_exit)
[文档]
def vtos(v):
"""
将给定的输入向量v转换为字符串。
参数:
- v (jittor.Var): 输入的需要转换的张量。
返回值:
output (str): 转换后的字符串。
代码示例:
>>> import jittor as jt
>>> v = jt.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
>>> print(jt.vtos(v)) # '1 2 3'
jt.Var([[1 2 3]
[1 2 3]
[1 2 3]], dtype=int32)
"""
data_str = f"jt.Var({v.numpy()}, dtype={v.dtype})"
data_str = data_str.replace("\n", "\n ")
return data_str
Var.__str__ = vtos
Var.__repr__ = vtos
Var.peek = lambda x: f"{x.dtype}{x.shape}"
[文档]
def size(v, dim=None):
"""
返回该张量的大小。 如果未指定dim，则返回值是该张量的大小。 如果指定了dim，则返回一个该维度的大小。
参数：
- v (jt.Var): 输入的张量。
- dim (int， 可选): 要查询大小的维度。默认值：``None``
返回值：
output{int, tuple}: 如果dim为None，返回v的总大小；否则返回v在维度dim上的大小。
代码示例:
>>> import jittor as jt
>>> a = jt.random((2,3))
>>> jt.size(a)
6
>>> jt.size(a, 0)
2
>>> jt.size(a, 1)
3
注意：
如果dim < 0，则将其解释为：v的维度数+dim，因此保持其在v的尺寸中。例如，如果v是2D的，dim= -1表示最后一个维度。
"""
if dim is None:
return v.shape
return v.shape[dim]
Var.size = size
[文档]
def to_int(v):
"""
将v中元素转化为整数.
参数:
- v (jittor.Var): 输入只能是单个元素的Var。
返回值:
output (int): 输出转化之后的整数。
代码示例:
>>> import jittor as jt
>>> print(jt.to_int(jt.Var(True)))
1
>>> print(jt.to_int(jt.Var([[1.6]])))
1
"""
return ori_int(v.item())
[文档]
def to_float(v):
"""将输入变量v中元素转化为浮点数类型.
参数:
- v (jittor.Var): 输入只能是单个元素的张量(jittor.Var)。
返回值:
output (float): 输出转化之后的浮点数。
代码示例:
>>> import jittor as jt
>>> print(jt.to_float(jt.Var(1)))
1.0
>>> print(jt.to_float(jt.Var([[True]])))
1.0
"""
return ori_float(v.item())
[文档]
def to_bool(v):
"""
将输入变量v中元素转化为布尔值.
参数:
- v (jittor.Var): 输入只能是单个元素的Var, 其中的元素必须是布尔值(bool)或者是整数(int)。
返回值:
output (bool): 输出转化之后的布尔值。
代码示例:
>>> import jittor as jt
>>> print(jt.to_bool(jt.Var([1])))
True
>>> print(jt.to_bool(jt.Var(0)))
False
>>> print(jt.to_bool(jt.Var(False)))
False
>>> print(jt.to_bool(jt.Var(9)))
True
"""
assert v.dtype.is_int() or v.dtype.is_bool()
return ori_bool(v.item())
Var.__int__ = to_int
Var.__float__ = to_float
Var.__bool__ = to_bool
Var.__format__ = format
[文档]
def get_len(var):
"""
获取输入张量的第一个维度的长度。
参数:
- var (jittor.Var): 输入的张量(jittor.Var)。
返回值:
output (int): 返回输入张量的第一个维度的长度。
代码示例:
>>> import jittor as jt
>>> v = jt.float32([1,2,3])
>>> len = jt.get_len(v)
>>> print(len)
3
注意:
此函数通过计算变量的shape属性来获取其长度，长度等于shape的第一个元素。
例如，如果变量的shape为(3, 2, 2)，则其长度为3。
"""
return var.shape[0]
Var.__len__ = get_len
int = int32
Var.int = Var.int32
Var.long = Var.int32
float = float32
Var.float = Var.float32
double = float64
Var.double = Var.float64
half = float16
Var.half = Var.float16
[文档]
def is_var(v):
"""
检查输入是否为Jittor的张量。
参数：
- v (Any): 待检查的输入对象。
返回值：
output (bool): 如果输入对象是Jittor变量，返回True，否则返回False。
代码示例:
>>> import jittor as jt
>>> a = jt.array([1, 2, 3])
>>> print(jt.is_var(a))
True
>>> b = 30
>>> print(jt.is_var(b))
False
"""
return isinstance(v, Var)
# __array__ interface is used for np.array(jt_var)
Var.__array__ = Var.numpy
Var.__array_priority__ = 2000
# __reduce__, __module__ is used for pickle.dump and pickle.load
Var.__module__ = "jittor"
Var.__reduce__ = lambda self: (Var, (self.data,))
from . import nn
from . import attention
from . import lr_scheduler
from . import linalg
from .linalg import einsum
from .nn import matmul, \
bmm, bmm_transpose, \
baddbmm
from . import contrib
from . import numpy2cupy
# from .contrib import concat, cat
from .misc import *
from . import sparse
from . import optim
from . import dataset
from . import init
dtype = NanoString
import jittor_utils
for backend in jittor_utils.backends:
if hasattr(backend, "post_process"):
backend.post_process()
# impl x.func(...) -> func_(...)
args = {"x", "input", "self"}
_white_list = {"mul", "add", "sub"}
for k,v in list(Var.__dict__.items()):
if k.startswith("_"): continue
if k.endswith("_"): continue
if not callable(v): continue
if k not in _white_list:
if not hasattr(v, "__code__"): continue
conames = v.__code__.co_varnames
if len(conames) == 0: continue
arg_name = conames[0]
if arg_name not in args: continue
new_k = k+"_"
if hasattr(Var, new_k): continue
def inplace_wrapper(new_k, prev_func):
setattr(Var, new_k, lambda x, *args, **kw: x.assign(prev_func(x, *args, **kw)))
inplace_wrapper(new_k, v)
from . import math_util
from .math_util import *
from . import distributions