tch-rs指南 - Tensor的基本 *** 作

tch-rs指南 - Tensor的基本 *** 作,第1张

文章目录
    • 1 概述
    • 2 Tensor的基本 *** 作
      • 2.1 Tensor的初始化
        • (1)通过数组创建
        • (2)通过默认方法创建
        • (3)通过其他的`tensor`创建
        • (4)通过`opencv::core::Mat`创建
      • 2.2 Tensor的属性
      • 2.3 Tensor的运算
        • (1)改变device
        • (2)获取值(indexing and slicing)
        • (3)合并tensors
        • (4)四则运算
    • 参考资料

1 概述

在使用rust进行torch模型部署时,不可避免地会用到tch-rs。但是tch-rs的文档太过简洁,和没有一样,网上的资料也少得可怜,很多 *** 作需要我们自己去试。这些内容虽然简单,但是自己找起来很费时间。

这篇文章总结了如何使用tch-rs进行tensor的基本 *** 作。讲述的内容参考了pytorch的tensor教程。

运行环境:

[dependencies]
tch = "0.7.0"
opencv = "0.63"
2 Tensor的基本 *** 作

用到的库

use std::iter;

use opencv::prelude::*;
use opencv::core::{Mat, Scalar};
use opencv::core::{CV_8UC3};
use tch::IndexOp;
use tch::{Device, Tensor};
2.1 Tensor的初始化 (1)通过数组创建
let t = Tensor::of_slice::<i32>(&[1, 2, 3, 4, 5]);
t.print();
// vector也是一样的
let v = vec![1,2,3];
let t = Tensor::of_slice::<i32>(&v);
t.print();
// 2d vector
let v = vec![[1.5,2.0,3.9,4.4], [3.1,4.3,5.1,6.9]];
let v:Vec<f32> = v
    .iter()
    .flat_map(|array| array.iter())
    .cloned()
    .collect();
let data = unsafe{
    std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * std::mem::size_of::<f32>())
};
let t = Tensor::of_data_size(data, &[2,4], tch::Kind::Float);
t.print();

print的结果是

 1
 2
 3
 4
 5
[ CPUIntType{5} ]
 1
 2
 3
[ CPUIntType{3} ]
 1.5000  2.0000  3.9000  4.4000
 3.1000  4.3000  5.1000  6.9000
[ CPUFloatType{2,4} ]
(2)通过默认方法创建
let t = Tensor::randn(&[2, 3], (tch::Kind::Float, Device::Cpu));
t.print();
let t = Tensor::ones(&[2, 3], (tch::Kind::Float, Device::Cpu));
t.print();
let t = Tensor::zeros(&[2, 3], (tch::Kind::Float, Device::Cpu));
t.print();
let t = Tensor::arange_start(0, 2 * 3, (tch::Kind::Float, Device::Cpu)).view([2, 3]);
t.print();

print的结果是

 1.0522  0.6981  0.9236
 0.2324 -1.1048 -2.5820
[ CPUFloatType{2,3} ]
 1  1  1
 1  1  1
[ CPUFloatType{2,3} ]
 0  0  0
 0  0  0
[ CPUFloatType{2,3} ]
 0  1  2
 3  4  5
[ CPUFloatType{2,3} ]
(3)通过其他的tensor创建
let t = Tensor::randn(&[2, 3], (tch::Kind::Float, Device::Cpu));
let t = t.rand_like();
t.print();

print的结果是

 0.3376  0.1885  0.3415
 0.5135  0.8321  0.4140
[ CPUFloatType{2,3} ]
(4)通过opencv::core::Mat创建

这可以用在opencv读取图像后,转为torch tensor。当然tch-rs本身也有各种读取图片的方式,可见tch::vision::image。这里介绍两种方法,一种通过tch::Tensor::f_of_blob,一种通过tch::Tensor::of_data_size

// 创建一个(row, col, channel)=(2, 3, 3)=(height, width, channel)的Mat
let mat = Mat::new_rows_cols_with_default(
    2, 3, CV_8UC3, Scalar::from((3.0, 2.0, 1.0))
).unwrap();
// 获取mat的size,这里的结果是[2, 3, 3]
let size: Vec<_> = mat.mat_size().iter().cloned().map(|dim| dim as i64).chain(iter::once(mat.channels() as i64)).collect();
// 获取每个dimension的stride,这里的结果是[9, 3, 1]
let strides = {
    let mut strides: Vec<_> = size
        .iter()
        .rev()
        .cloned()
        .scan(1, |prev, dim| {
            let stride = *prev;
            *prev *= dim;
            Some(stride)
        })
        .collect();
    strides.reverse();
    strides
};
// 构建tensor
let t = unsafe {
    let ptr = mat.ptr(0).unwrap() as *const u8;
    tch::Tensor::f_of_blob(ptr, &size, &strides, tch::Kind::Uint8, tch::Device::Cpu).unwrap()
};
t.print();

print的结果是

(1,.,.) = 
  3  2  1
  3  2  1
  3  2  1

(2,.,.) = 
  3  2  1
  3  2  1
  3  2  1
[ CPUByteType{2,3,3} ]

还有一种比较简洁的转换方法

let mut mat = Mat::new_rows_cols_with_default(
    2, 3, CV_8UC3, Scalar::from((3.0, 2.0, 1.0))
).unwrap();
let h = mat.size().unwrap().height;
let w = mat.size().unwrap().width;   
let data = mat.data_bytes_mut().unwrap(); 
let t = tch::Tensor::of_data_size(data, &[h as i64, w as i64, 3], tch::Kind::Uint8);
t.print();

print的结果也是

(1,.,.) = 
  3  2  1
  3  2  1
  3  2  1

(2,.,.) = 
  3  2  1
  3  2  1
  3  2  1
[ CPUByteType{2,3,3} ]
test tensor_ops::init_ops ... ok
2.2 Tensor的属性

用tch::Tensor的print()方法可打印出数据的所有属性,但是想要获取到这些属性,需要用其他的方法。

let t = Tensor::randn(&[2, 3], (tch::Kind::Float, Device::Cpu));
println!("size of the tensor: {:?}", t.size());
println!("kind of the tensor: {:?}", t.kind());
println!("device on which the tensor is located: {:?}", t.device());

打印的结果是

size of the tensor: [2, 3]
kind of the tensor: Float
device on which the tensor is located: Cpu
2.3 Tensor的运算 (1)改变device

.to().to_device()这两个方法都可以。

let mut t = Tensor::randn(&[2, 3], (tch::Kind::Float, Device::Cpu));
if tch::Cuda::is_available(){
    t = t.to(Device::Cuda(0));
    println!("change device to {:?}", t.device());
}
t = t.to_device(Device::Cpu);
println!("change device to {:?}", t.device());

如果是有cuda,且安装了cuda版本的tch-rs的话,就会打印出

change device to Cuda(0)
change device to Cpu
(2)获取值(indexing and slicing)

这个在tch-rs的例子中有很多,详见tests/tensor_indexing.rs。这里列几种常用的。

通过.i()进行索引

let tensor = Tensor::arange_start(0, 2 * 3, (tch::Kind::Float, Device::Cpu)).view([2, 3]);
println!("original tensor:");
tensor.print();
println!("tensor.i(0):");
tensor.i(0).print();
println!("tensor.i((1, 1)):");
tensor.i((1, 1)).print();
println!("tensor.i((.., 2)):");
tensor.i((.., 2)).print();
println!("tensor.i((.., -1)):");
tensor.i((.., -1)).print();
println!("tensor.i((.., [2, 0])):");
let index: &[_] = &[2, 0];
tensor.i((.., index)).print();

打印的结果是

original tensor:
 0  1  2
 3  4  5
[ CPUFloatType{2,3} ]
tensor.i(0):
 0
 1
 2
[ CPUFloatType{3} ]
tensor.i((1, 1)):
4
[ CPUFloatType{} ]
tensor.i((.., 2)):
 2
 5
[ CPUFloatType{2} ]
tensor.i((.., -1)):
 2
 5
[ CPUFloatType{2} ]
tensor.i((.., [2, 0])):
 2  0
 5  3
[ CPUFloatType{2,2} ]

通过.index()进行索引

let tensor = Tensor::arange(6, (tch::Kind::Int64, Device::Cpu)).view((2, 3));
println!("original tensor:");
tensor.print();
let rows_select = Tensor::of_slice(&[0i64, 1, 0]);
let column_select = Tensor::of_slice(&[1i64, 2, 2]);
let selected = tensor.index(&[Some(rows_select), Some(column_select)]);
println!("selecte by row and column:");
selected.print();

打印的结果是

original tensor:
 0  1  2
 3  4  5
[ CPULongType{2,3} ]
selecte by row and column:
 1
 5
 2
[ CPULongType{3} ]
(3)合并tensors

Tensor::f_cat不会生成新的axis,而Tensor::stack会生成新的axis。

let t1 = Tensor::arange(6, (tch::Kind::Int64, Device::Cpu)).view((2, 3));
let t2 = Tensor::arange_start(6, 12, (tch::Kind::Int64, Device::Cpu)).view((2, 3));
let tensor = Tensor::f_cat(&[t1.copy(), t2.copy()], 1).unwrap();
println!("using Tensor::f_cat");
tensor.print();
let tensor = Tensor::stack(&[t1.copy(), t2.copy()], 1);
println!("using Tensor::stack");
tensor.print();

打印的结果是

using Tensor::f_cat
  0   1   2   6   7   8
  3   4   5   9  10  11
[ CPULongType{2,6} ]
using Tensor::stack
(1,.,.) = 
  0  1  2
  6  7  8

(2,.,.) = 
   3   4   5
   9  10  11
[ CPULongType{2,2,3} ]
(4)四则运算

tch-rs对[+, -, *, /]都进行了重载,可以实现和标量的直接运算。涉及到dim的复杂运算可以用tensor来处理。下面以加法为例,其他与f_add对应的分别是f_subf_mulf_div

let tensor = Tensor::ones(&[2, 4, 3], (tch::Kind::Float, Device::Cpu));
tensor.print();
// add with scalar
let add_tensor = &tensor + 0.5;
add_tensor.print();
// add with tensor
let add_tensor = Tensor::of_slice::<f32>(&[1.0,2.0,3.0]).view((1,1,3));
let add_tensor = &tensor.f_add(&add_tensor).unwrap();
add_tensor.print();

打印的结果为

original tensor:
(1,.,.) = 
  1  1  1
  1  1  1
  1  1  1
  1  1  1

(2,.,.) = 
  1  1  1
  1  1  1
  1  1  1
  1  1  1
[ CPUFloatType{2,4,3} ]
add with scalar:
(1,.,.) = 
  1.5000  1.5000  1.5000
  1.5000  1.5000  1.5000
  1.5000  1.5000  1.5000
  1.5000  1.5000  1.5000

(2,.,.) = 
  1.5000  1.5000  1.5000
  1.5000  1.5000  1.5000
  1.5000  1.5000  1.5000
  1.5000  1.5000  1.5000
[ CPUFloatType{2,4,3} ]
add with tensor:
(1,.,.) = 
  2  3  4
  2  3  4
  2  3  4
  2  3  4

(2,.,.) = 
  2  3  4
  2  3  4
  2  3  4
  2  3  4
[ CPUFloatType{2,4,3} ]
参考资料

[1] https://github.com/LaurentMazare/tch-rs
[2] https://pytorch.org/tutorials/beginner/basics/tensorqs_tutorial.html#

欢迎分享,转载请注明来源:内存溢出

原文地址: http://outofmemory.cn/langs/915546.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-05-16
下一篇 2022-05-16

发表评论

登录后才能评论

评论列表(0条)

保存