← Back

Module nn

trait Module

Trait for all neural network modules

Source: nn.joule:9

fn forward(&self, input: &Tensor) -> Tensor;

Forward pass

Source: nn.joule:11

fn parameters(&self) -> Vec<&Tensor>;

Get all trainable parameters

Source: nn.joule:14

fn parameters_mut(&mut self) -> Vec<&mut Tensor>;

Get mutable references to parameters

Source: nn.joule:17

fn train(&mut self, mode: bool);

Set training mode

Source: nn.joule:20

fn training(&self) -> bool;

Check if in training mode

Source: nn.joule:23

fn to(&mut self, device: Device);

Move module to device

Source: nn.joule:26

fn num_parameters(&self) -> usize

Number of parameters

Source: nn.joule:29

fn zero_grad(&mut self)

Zero gradients

Source: nn.joule:34

struct Linear

Linear (fully connected) layer: y = xW^T + b

Source: nn.joule:46

fn new(in_features: usize, out_features: usize, bias: bool) -> Self

Source: nn.joule:57

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:81

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:90

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:98

fn train(&mut self, mode: bool)

Source: nn.joule:106

fn training(&self) -> bool

Source: nn.joule:110

fn to(&mut self, device: Device)

Source: nn.joule:114

struct Bilinear

Bilinear layer: y = x1 A x2 + b

Source: nn.joule:123

fn new(

Source: nn.joule:133

fn forward(&self, x1: &Tensor, x2: &Tensor) -> Tensor

Source: nn.joule:160

struct Conv1d

1D Convolution

Source: nn.joule:184

fn new(

Source: nn.joule:198

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:235

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:249

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:257

fn train(&mut self, mode: bool)

Source: nn.joule:265

fn training(&self) -> bool

Source: nn.joule:266

fn to(&mut self, device: Device)

Source: nn.joule:268

struct Conv2d

2D Convolution

Source: nn.joule:277

fn new(

Source: nn.joule:291

fn simple(

Create with simple parameters (kernel_size, stride, padding as single values)

Source: nn.joule:329

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:352

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:366

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:372

fn train(&mut self, mode: bool)

Source: nn.joule:378

fn training(&self) -> bool

Source: nn.joule:379

fn to(&mut self, device: Device)

Source: nn.joule:381

struct ConvTranspose2d

Transposed 2D Convolution (for upsampling)

Source: nn.joule:388

fn new(

Source: nn.joule:403

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:445

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:460

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:466

fn train(&mut self, mode: bool)

Source: nn.joule:472

fn training(&self) -> bool

Source: nn.joule:473

fn to(&mut self, device: Device)

Source: nn.joule:475

struct BatchNorm1d

Batch Normalization

Source: nn.joule:486

fn new(num_features: usize, eps: f64, momentum: f64) -> Self

Source: nn.joule:498

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:514

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:543

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:547

fn train(&mut self, mode: bool)

Source: nn.joule:551

fn training(&self) -> bool

Source: nn.joule:552

fn to(&mut self, device: Device)

Source: nn.joule:554

struct BatchNorm2d

Batch Normalization 2D

Source: nn.joule:563

fn new(num_features: usize, eps: f64, momentum: f64) -> Self

Source: nn.joule:575

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:591

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:619

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:623

fn train(&mut self, mode: bool)

Source: nn.joule:627

fn training(&self) -> bool

Source: nn.joule:628

fn to(&mut self, device: Device)

Source: nn.joule:630

struct LayerNorm

Layer Normalization

Source: nn.joule:639

fn new(normalized_shape: &[usize], eps: f64) -> Self

Source: nn.joule:648

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:661

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:677

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:681

fn train(&mut self, mode: bool)

Source: nn.joule:685

fn training(&self) -> bool

Source: nn.joule:686

fn to(&mut self, device: Device)

Source: nn.joule:688

struct GroupNorm

Group Normalization

Source: nn.joule:695

fn new(num_groups: usize, num_channels: usize, eps: f64) -> Self

Source: nn.joule:705

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:721

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:753

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:757

fn train(&mut self, mode: bool)

Source: nn.joule:761

fn training(&self) -> bool

Source: nn.joule:762

fn to(&mut self, device: Device)

Source: nn.joule:764

struct ReLU;

ReLU activation

Source: nn.joule:775

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:779

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:783

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:784

fn train(&mut self, _mode: bool)

Source: nn.joule:785

fn training(&self) -> bool

Source: nn.joule:786

fn to(&mut self, _device: Device)

Source: nn.joule:787

struct LeakyReLU

Leaky ReLU activation

Source: nn.joule:791

fn new(negative_slope: f64) -> Self

Source: nn.joule:796

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:803

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:807

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:808

fn train(&mut self, _mode: bool)

Source: nn.joule:809

fn training(&self) -> bool

Source: nn.joule:810

fn to(&mut self, _device: Device)

Source: nn.joule:811

struct GELU;

GELU activation

Source: nn.joule:815

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:819

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:823

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:824

fn train(&mut self, _mode: bool)

Source: nn.joule:825

fn training(&self) -> bool

Source: nn.joule:826

fn to(&mut self, _device: Device)

Source: nn.joule:827

struct SiLU;

SiLU (Swish) activation

Source: nn.joule:831

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:835

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:839

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:840

fn train(&mut self, _mode: bool)

Source: nn.joule:841

fn training(&self) -> bool

Source: nn.joule:842

fn to(&mut self, _device: Device)

Source: nn.joule:843

struct Sigmoid;

Sigmoid activation

Source: nn.joule:847

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:851

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:855

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:856

fn train(&mut self, _mode: bool)

Source: nn.joule:857

fn training(&self) -> bool

Source: nn.joule:858

fn to(&mut self, _device: Device)

Source: nn.joule:859

struct Tanh;

Tanh activation

Source: nn.joule:863

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:867

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:871

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:872

fn train(&mut self, _mode: bool)

Source: nn.joule:873

fn training(&self) -> bool

Source: nn.joule:874

fn to(&mut self, _device: Device)

Source: nn.joule:875

struct Softmax

Softmax activation

Source: nn.joule:879

fn new(dim: i64) -> Self

Source: nn.joule:884

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:891

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:895

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:896

fn train(&mut self, _mode: bool)

Source: nn.joule:897

fn training(&self) -> bool

Source: nn.joule:898

fn to(&mut self, _device: Device)

Source: nn.joule:899

struct MaxPool2d

Max Pooling 2D

Source: nn.joule:907

fn new(

Source: nn.joule:914

fn simple(kernel_size: usize, stride: Option<usize>) -> Self

Source: nn.joule:926

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:934

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:938

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:939

fn train(&mut self, _mode: bool)

Source: nn.joule:940

fn training(&self) -> bool

Source: nn.joule:941

fn to(&mut self, _device: Device)

Source: nn.joule:942

struct AvgPool2d

Average Pooling 2D

Source: nn.joule:946

fn new(

Source: nn.joule:953

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:968

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:972

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:973

fn train(&mut self, _mode: bool)

Source: nn.joule:974

fn training(&self) -> bool

Source: nn.joule:975

fn to(&mut self, _device: Device)

Source: nn.joule:976

struct AdaptiveAvgPool2d

Adaptive Average Pooling 2D

Source: nn.joule:980

fn new(output_size: (usize, usize)) -> Self

Source: nn.joule:985

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:992

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:996

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:997

fn train(&mut self, _mode: bool)

Source: nn.joule:998

fn training(&self) -> bool

Source: nn.joule:999

fn to(&mut self, _device: Device)

Source: nn.joule:1000

struct Dropout

Dropout

Source: nn.joule:1008

fn new(p: f64) -> Self

Source: nn.joule:1014

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:1021

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:1029

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:1030

fn train(&mut self, mode: bool)

Source: nn.joule:1031

fn training(&self) -> bool

Source: nn.joule:1032

fn to(&mut self, _device: Device)

Source: nn.joule:1033

struct Dropout2d

Dropout 2D (spatial dropout for CNNs)

Source: nn.joule:1037

fn new(p: f64) -> Self

Source: nn.joule:1043

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:1050

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:1059

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:1060

fn train(&mut self, mode: bool)

Source: nn.joule:1061

fn training(&self) -> bool

Source: nn.joule:1062

fn to(&mut self, _device: Device)

Source: nn.joule:1063

struct LSTMCell

LSTM cell

Source: nn.joule:1071

fn new(input_size: usize, hidden_size: usize, bias: bool) -> Self

Source: nn.joule:1082

fn forward_step(

Source: nn.joule:1106

struct LSTM

Full LSTM layer

Source: nn.joule:1131

fn new(

Source: nn.joule:1141

fn forward(

Source: nn.joule:1176

struct MultiheadAttention

Multi-Head Attention

Source: nn.joule:1251

fn new(embed_dim: usize, num_heads: usize, dropout: f64) -> Self

Source: nn.joule:1264

fn forward(

Source: nn.joule:1284

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:1339

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:1344

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:1353

fn train(&mut self, mode: bool)

Source: nn.joule:1362

fn training(&self) -> bool

Source: nn.joule:1370

fn to(&mut self, device: Device)

Source: nn.joule:1372

struct Embedding

Embedding lookup table

Source: nn.joule:1385

fn new(num_embeddings: usize, embedding_dim: usize, padding_idx: Option<usize>) -> Self

Source: nn.joule:1394

fn from_pretrained(weight: Tensor, freeze: bool, padding_idx: Option<usize>) -> Self

Create from pretrained weights

Source: nn.joule:1407

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:1427

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:1431

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:1435

fn train(&mut self, mode: bool)

Source: nn.joule:1439

fn training(&self) -> bool

Source: nn.joule:1440

fn to(&mut self, device: Device)

Source: nn.joule:1442

struct Sequential

Sequential container

Source: nn.joule:1452

fn new() -> Self

Source: nn.joule:1458

fn add<M: Module + 'static>(&mut self, module: M)

Source: nn.joule:1465

fn forward(&self, input: &Tensor) -> Tensor

Source: nn.joule:1471

fn parameters(&self) -> Vec<&Tensor>

Source: nn.joule:1479

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Source: nn.joule:1483

fn train(&mut self, mode: bool)

Source: nn.joule:1487

fn training(&self) -> bool

Source: nn.joule:1494

fn to(&mut self, device: Device)

Source: nn.joule:1496

fn linear(input: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor

Source: nn.joule:1512

fn conv2d(

Source: nn.joule:1522

fn batch_norm(

Source: nn.joule:1540

fn dropout(input: &Tensor, p: f64, training: bool) -> Tensor

Source: nn.joule:1580

fn max_pool2d(

Source: nn.joule:1590

fn softmax(input: &Tensor, dim: i64) -> Tensor

Source: nn.joule:1601

fn log_softmax(input: &Tensor, dim: i64) -> Tensor

Source: nn.joule:1607

fn relu(input: &Tensor) -> Tensor

Source: nn.joule:1613

fn gelu(input: &Tensor) -> Tensor

Source: nn.joule:1619

fn sigmoid(input: &Tensor) -> Tensor

Source: nn.joule:1625

fn tanh(input: &Tensor) -> Tensor

Source: nn.joule:1631

fn cross_entropy(

Source: nn.joule:1637

fn binary_cross_entropy_with_logits(

Source: nn.joule:1668

fn mse_loss(input: &Tensor, target: &Tensor, reduction: &str) -> Tensor

Source: nn.joule:1695

fn l1_loss(input: &Tensor, target: &Tensor, reduction: &str) -> Tensor

Source: nn.joule:1708