温馨提示×

Torch中Module怎么定义一个神经网络结构

小亿
82
2024-03-19 13:26:37
栏目: 深度学习

在Torch中,可以通过继承nn.Module类来定义一个神经网络结构。以下是一个简单的示例:

import torch
import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNN()

在这个示例中,SimpleNN类继承自nn.Module,并在__init__方法中定义了神经网络的结构,包括两个全连接层和一个ReLU激活函数。在forward方法中定义了数据在神经网络中的流动,也就是前向传播过程。最后通过实例化SimpleNN类来创建一个神经网络模型。

0