Struct ConvShape

Source

pub struct ConvShape {
    pub c: usize,
    pub h: usize,
    pub w: usize,
    pub o: usize,
    pub kh: usize,
    pub kw: usize,
    pub stride: (usize, usize),
    pub pad: (usize, usize),
    pub dilation: (usize, usize),
    pub groups: usize,
}

Expand description

一次卷积的全部尺寸，与 pagesift_infer::ops::conv2d 的参数一一对应：输入 [c, h, w]，权重 [o, c/groups, kh, kw]。

Fields§

§c: usize§h: usize§w: usize§o: usize§kh: usize§kw: usize§stride: (usize, usize)§pad: (usize, usize)§dilation: (usize, usize)§groups: usize

Implementations§

Source §

impl ConvShape

Source

pub fn out_hw(&self) -> (usize, usize)

输出的 (OH, OW)，公式同 PyTorch。

Examples found in repository ?

examples/gpu_bench.rs (line 166)

62fn main() {
63    let g = match Gpu::open() {
64        Ok(g) => g,
65        Err(e) => {
66            println!("没有可用的 GPU：{e}");
67            return;
68        }
69    };
70    println!("GPU：{}（算力 {:?}）\n", g.name(), g.compute_capability());
71    println!(
72        "| 项目 | CPU ms | GPU 内核 ms | GPU 端到端 ms | 倍数（内核） | 倍数（端到端） | CPU GFLOPS | GPU GFLOPS |"
73    );
74    println!("|---|---:|---:|---:|---:|---:|---:|---:|");
75    let mut rng = Rng(0x9e37_79b9_7f4a_7c15);
76
77    for s in [512, 1024, 2048] {
78        let a = rng.tensor(vec![s, s]);
79        let b = rng.tensor(vec![s, s]);
80        let gflop = 2.0 * (s as f64).powi(3) / 1e9;
81        let cpu = wall_ms(3, || drop(ops::matmul(&a, &b)));
82        let (da, db) = (g.upload(&a.data).unwrap(), g.upload(&b.data).unwrap());
83        let mut dc = g.alloc(s * s).unwrap();
84        let kern = gpu_ms(&g, 20, || {
85            g.gemm_device(&da, &db, &mut dc, s, s, s).unwrap()
86        });
87        let e2e = wall_ms(5, || drop(g.matmul(&a, &b).unwrap()));
88        row(&format!("SGEMM {s}³"), gflop, cpu, kern, e2e);
89    }
90
91    // (说明, [C, H, W], [O, kh, kw], stride, pad, groups)
92    type Conv = (
93        &'static str,
94        [usize; 3],
95        [usize; 3],
96        (usize, usize),
97        (usize, usize),
98        usize,
99    );
100    let convs: [Conv; 6] = [
101        (
102            "3×3 s2 首层 3→32 @640²",
103            [3, 640, 640],
104            [32, 3, 3],
105            (2, 2),
106            (1, 1),
107            1,
108        ),
109        (
110            "1×1 64→128 @160²",
111            [64, 160, 160],
112            [128, 1, 1],
113            (1, 1),
114            (0, 0),
115            1,
116        ),
117        (
118            "3×3 64→64 @80²",
119            [64, 80, 80],
120            [64, 3, 3],
121            (1, 1),
122            (1, 1),
123            1,
124        ),
125        (
126            "3×3 s(2,1) 64→128 @48×320",
127            [64, 48, 320],
128            [128, 3, 3],
129            (2, 1),
130            (1, 1),
131            1,
132        ),
133        (
134            "深度 3×3 128 @160²",
135            [128, 160, 160],
136            [128, 3, 3],
137            (1, 1),
138            (1, 1),
139            128,
140        ),
141        (
142            "深度 7×7 96 @80²",
143            [96, 80, 80],
144            [96, 7, 7],
145            (1, 1),
146            (3, 3),
147            96,
148        ),
149    ];
150    for (name, [c, h, w], [o, kh, kw], stride, pad, groups) in convs {
151        let x = rng.tensor(vec![c, h, w]);
152        let wt = rng.tensor(vec![o, c / groups, kh, kw]);
153        let bias: Vec<f32> = rng.tensor(vec![o]).data;
154        let s = ConvShape {
155            c,
156            h,
157            w,
158            o,
159            kh,
160            kw,
161            stride,
162            pad,
163            dilation: (1, 1),
164            groups,
165        };
166        let (oh, ow) = s.out_hw();
167        let gflop = 2.0 * (o * oh * ow * (c / groups) * kh * kw) as f64 / 1e9;
168        let cpu = wall_ms(3, || {
169            drop(ops::conv2d(
170                &x,
171                &wt,
172                Some(&bias),
173                stride,
174                pad,
175                (1, 1),
176                groups,
177            ))
178        });
179        let (dx, dw, db) = (
180            g.upload(&x.data).unwrap(),
181            g.upload(&wt.data).unwrap(),
182            g.upload(&bias).unwrap(),
183        );
184        let mut out = g.alloc(o * oh * ow).unwrap();
185        let kern = gpu_ms(&g, 20, || {
186            g.conv2d_device_into(&dx, &dw, Some(&db), &s, &mut out)
187                .unwrap()
188        });
189        let e2e = wall_ms(5, || {
190            drop(
191                g.conv2d(&x, &wt, Some(&bias), stride, pad, (1, 1), groups)
192                    .unwrap(),
193            )
194        });
195        row(name, gflop, cpu, kern, e2e);
196    }
197}