PyTorch/CSPRNG
torchcsprng 是一个 PyTorch C++/CUDA 扩展,提供以下功能:
- AES 128位加密/解密,支持两种模式:ECB 和 CTR
- 为 PyTorch 提供密码学安全的伪随机数生成器。
设计
torchcsprng 在 CPU 上使用其生成器之一生成随机的 128 位密钥,并在 CPU 或 GPU 上使用 CUDA 运行 CTR 模式 的 AES128,以生成随机的 128 位状态,并应用转换函数将其映射到目标张量值。 这种方法基于 Parallel Random Numbers: As Easy as 1, 2, 3(John K. Salmon, Mark A. Moraes, Ron O. Dror, and David E. Shaw, D. E. Shaw Research)。 这使得 torchcsprng 在 CUDA 和 CPU 上既具有密码学安全性,又可以并行运行。
优势:
- 用户可以选择基于种子(用于测试)或基于随机设备(完全密码学安全)的生成器
- 一个生成器实例可同时用于 CPU 和 CUDA 张量(因为加密密钥始终在 CPU 上生成)
- CPU 随机数生成也是并行的(与默认的 PyTorch CPU 生成器不同)
特性
torchcsprng 0.2.0 公开了新的张量加密/解密 API。张量加密/解密 API 与数据类型无关,因此任何数据类型的张量都可以被加密,结果可以存储到任何数据类型的张量中。加密密钥也可以是任何数据类型的张量。目前 torchcsprng 支持两种模式的 128 位密钥 AES 加密:ECB 和 CTR。
torchcsprng.encrypt(input: Tensor, output: Tensor, key: Tensor, cipher: string, mode: string)
torchcsprng.decrypt(input: Tensor, output: Tensor, key: Tensor, cipher: string, mode: string)
torchcsprng 提供两种方法来创建密码学安全和非密码学安全的 PRNG:
创建 PRNG 的方法 | 是否密码学安全? | 是否有种子? | 底层实现 |
---|---|---|---|
create_random_device_generator(token: string=None) | 是 | 否 | 参见 std::random_device 及其构造函数。libstdc++ 中的实现期望 token 命名随机字节的来源。可能的 token 值包括 "default"、"rand_s"、"rdseed"、"rdrand"、"rdrnd"、"/dev/urandom"、"/dev/random"、"mt19937",以及指定 mt19937 引擎种子的整数字符串。(除 "default" 外的 token 值仅对某些目标有效。)如果 token=None,则使用实现定义的 token 构造新的 std::random_device 对象。 |
create_mt19937_generator(seed: int=None) | 否 | 是 | 参见 std::mt19937 及其构造函数。构造一个 mersenne_twister_engine 对象,并将其内部状态序列初始化为伪随机值。如果 seed=None,则使用 default_seed 为引擎设置种子。 |
以下方法列表支持所有上述 PRNG:
核心函数 | CUDA | CPU |
---|---|---|
random_() | 是 | 是 |
random_(to) | 是 | 是 |
random_(from, to) | 是 | 是 |
uniform_(from, to) | 是 | 是 |
normal_(mean, std) | 是 | 是 |
cauchy_(median, sigma) | 是 | 是 |
log_normal_(mean, std) | 是 | 是 |
geometric_(p) | 是 | 是 |
exponential_(lambda) | 是 | 是 |
randperm(n) | 是* | 是 |
- 计算在 CPU 上完成,结果复制到 CUDA
安装
CSPRNG 支持 Python 3.6-3.9,可在以下操作系统上运行,并可用于以下设备上的 PyTorch 张量:
张量设备类型 | Linux | macOS | MS Windows |
---|---|---|---|
CPU | 支持 | 支持 | 支持 |
CUDA | 支持 | 不支持 | 自 0.2.0 起支持 |
以下是对应的 CSPRNG 版本和支持的 Python 版本。
PyTorch | CSPRNG | Python | CUDA |
---|---|---|---|
1.8.0 | 0.2.0 | 3.7-3.9 | 10.1, 10.2, 11.1 |
1.7.1 | 0.1.4 | 3.6-3.8 | 9.2, 10.1, 10.2 |
1.7.0 | 0.1.3 | 3.6-3.8 | 9.2, 10.1, 10.2 |
1.6.0 | 0.1.2 | 3.6-3.8 | 9.2, 10.1, 10.2 |
二进制安装
Anaconda:
操作系统 | CUDA | |
---|---|---|
Linux/Windows | 10.1 10.2 11.1 None | conda install torchcsprng cudatoolkit=10.1 -c pytorch -c conda-forge conda install torchcsprng cudatoolkit=10.2 -c pytorch -c conda-forge conda install torchcsprng cudatoolkit=11.1 -c pytorch -c conda-forge conda install torchcsprng cpuonly -c pytorch -c conda-forge |
macOS | None | conda install torchcsprng -c pytorch |
pip:
操作系统 | CUDA | |
---|---|---|
Linux/Windows | 10.1 10.2 11.1 None | pip install torchcsprng==0.2.0+cu101 torch==1.8.0+cu101 -f https://download.pytorch.org/whl/cu101/torch_stable.html pip install torchcsprng==0.2.0 torch==1.8.0 -f https://download.pytorch.org/whl/cu102/torch_stable.html pip install torchcsprng==0.2.0+cu111 torch==1.8.0+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html pip install torchcsprng==0.2.0+cpu torch==1.8.0+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html |
macOS | None | pip install torchcsprng torch |
每日构建版本:
Anaconda:
操作系统 | CUDA | |
---|---|---|
Linux/Windows | 10.1 10.2 11.1 None | conda install torchcsprng cudatoolkit=10.1 -c pytorch-nightly -c conda-forge conda install torchcsprng cudatoolkit=10.2 -c pytorch-nightly -c conda-forge conda install torchcsprng cudatoolkit=11.1 -c pytorch-nightly -c conda-forge conda install torchcsprng cpuonly -c pytorch-nightly -c conda-forge |
macOS | None | conda install torchcsprng -c pytorch-nightly |
pip:
操作系统 | CUDA | |
---|---|---|
Linux/Windows | 10.1 10.2 11.1 无 | pip install --pre torchcsprng -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html pip install --pre torchcsprng -f https://download.pytorch.org/whl/nightly/cu102/torch_nightly.html pip install --pre torchcsprng -f https://download.pytorch.org/whl/nightly/cu111/torch_nightly.html pip install --pre torchcsprng -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html |
macOS | 无 | pip install --pre torchcsprng -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html |
从源码安装
torchcsprng 是一个依赖于 PyTorch 的 Python C++/CUDA 扩展。为了从源码构建 CSPRNG,需要安装 Python(>=3.7) 和 PyTorch(>=1.8.0),以及 C++ 编译器(Linux 上的 gcc/clang,macOS 上的 XCode,Windows 上的 Visual Studio)。 要构建 torchcsprng,可以运行以下命令:
python setup.py install
默认情况下,如果找到 CUDA 并且 torch.cuda.is_available() 为 True,则会构建 GPU 支持。此外,可以通过设置 FORCE_CUDA=1 环境变量来强制构建 GPU 支持,这在构建 docker 镜像时很有用。
入门
torchcsprng API 可在 torchcsprng
模块中使用:
import torch
import torchcsprng as csprng
从 /dev/urandom 创建加密安全的 PRNG:
urandom_gen = csprng.create_random_device_generator('/dev/urandom')
在 CUDA 上创建空的布尔张量,并使用 urandom_gen 的随机值初始化:
torch.empty(10, dtype=torch.bool, device='cuda').random_(generator=urandom_gen)
tensor([ True, False, False, True, False, False, False, True, False, False],
device='cuda:0')
在 CUDA 上创建空的 int16 张量,并使用 urandom_gen 的范围在 [0, 100) 内的随机值初始化:
torch.empty(10, dtype=torch.int16, device='cuda').random_(100, generator=urandom_gen)
tensor([59, 20, 68, 51, 18, 37, 7, 54, 74, 85], device='cuda:0',
dtype=torch.int16)
创建非加密安全的 MT19937 PRNG:
mt19937_gen = csprng.create_mt19937_generator()
torch.empty(10, dtype=torch.int64, device='cuda').random_(torch.iinfo(torch.int64).min, to=None, generator=mt19937_gen)
tensor([-7584783661268263470, 2477984957619728163, -3472586837228887516,
-5174704429717287072, 4125764479102447192, -4763846282056057972,
-182922600982469112, -498242863868415842, 728545841957750221,
7740902737283645074], device='cuda:0')
从默认随机设备创建加密安全的 PRNG:
default_device_gen = csprng.create_random_device_generator()
torch.randn(10, device='cuda', generator=default_device_gen)
tensor([ 1.2885, 0.3240, -1.1813, 0.8629, 0.5714, 2.3720, -0.5627, -0.5551,
-0.6304, 0.1090], device='cuda:0')
使用种子创建非加密安全的 MT19937 PRNG:
mt19937_gen = csprng.create_mt19937_generator(42)
torch.empty(10, device='cuda').geometric_(p=0.2, generator=mt19937_gen)
tensor([ 7., 1., 8., 1., 11., 3., 1., 1., 5., 10.], device='cuda:0')
使用相同的种子重新创建 MT19937 PRNG:
mt19937_gen = csprng.create_mt19937_generator(42)
torch.empty(10, device='cuda').geometric_(p=0.2, generator=mt19937_gen)
tensor([ 7., 1., 8., 1., 11., 3., 1., 1., 5., 10.], device='cuda:0')
贡献
我们感谢所有的贡献。如果您计划贡献错误修复,请直接进行而无需进一步讨论。如果您计划贡献新功能、实用函数或扩展,请先开启一个 issue 并与我们讨论该功能。
许可证
torchcsprng 使用 BSD 3-clause 许可证。许可证文件可在此处查看。
使用条款
隐私政策
版权所有 © 2020 Meta Platforms, Inc