Implementation:LaurentMazare Tch rs Nn Conv2d

Knowledge Sources	tch-rs
Domains	Computer_Vision, Neural_Network_Layers
Last Updated	2026-02-08 14:00 GMT

Overview

Concrete tool for creating 2D convolutional layers provided by the tch nn module.

Description

nn::conv2d creates a Conv2D struct with a weight tensor of shape [o, i/groups, k, k] and optional bias of shape [o]. Parameters are registered in the VarStore. The Conv2D implements the Module trait and applies 2D convolution in its forward method, supporting configurable stride, padding, dilation, groups, and padding mode.

Usage

Use this to build convolutional neural networks for image processing tasks. Configure stride, padding, and kernel size via ConvConfig.

Code Reference

Source Location

Repository: tch-rs
File: src/nn/conv.rs
Lines: 184-186

Signature

pub fn conv2d<'a, T: Borrow<Path<'a>>>(
    vs: T,
    i: i64,
    o: i64,
    k: i64,
    c: ConvConfig,
) -> Conv2D

Import

use tch::nn;

I/O Contract

Inputs

Name	Type	Required	Description
vs	T: Borrow<Path>	Yes	VarStore path for parameter registration
i	i64	Yes	Number of input channels
o	i64	Yes	Number of output channels
k	i64	Yes	Square kernel size
c	ConvConfig	Yes	Config: stride, padding, dilation, groups, bias, padding_mode

Outputs

Name	Type	Description
Conv2D	nn::Conv2D	Struct with weight [o, i/groups, k, k], optional bias [o], implementing Module

Usage Examples

MNIST CNN Layer

use tch::{nn, nn::Module, Device};

let vs = nn::VarStore::new(Device::Cpu);
let conv1 = nn::conv2d(vs.root() / "conv1", 1, 32, 5, Default::default());

// Input: [batch, 1, 28, 28]
let input = tch::Tensor::randn([1, 1, 28, 28], tch::kind::FLOAT_CPU);
let output = conv1.forward(&input);  // shape: [1, 32, 24, 24]

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment