nn.Module类特性的代码
时间: 2023-10-23 15:28:45 浏览: 36
以下是使用 `nn.Module` 类的一些特性的示例代码:
1. 自动求导:
```python
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(3, 10)
self.fc2 = nn.Linear(10, 1)
def forward(self, x):
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
net = Net()
input = torch.randn(1, 3)
output = net(input)
output.backward()
```
在上述代码中,定义了一个名为 `Net` 的神经网络模型,继承自 `nn.Module`。在模型的前向传播过程中,使用了 PyTorch 提供的自动求导功能计算了输入 `input` 对输出 `output` 的梯度,并通过 `output.backward()` 将梯度反向传播到模型中的参数中。
2. 参数管理:
```python
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(3, 10)
self.fc2 = nn.Linear(10, 1)
def forward(self, x):
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
net = Net()
params = list(net.parameters())
print(params)
```
在上述代码中,定义了一个名为 `Net` 的神经网络模型,继承自 `nn.Module`。通过 `net.parameters()` 方法可以获取模型中的所有参数,包括全连接层 `fc1` 和 `fc2` 中的权重和偏置。
3. 组合模型:
```python
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.pool = nn.MaxPool2d(2)
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = self.conv1(x)
x = nn.functional.relu(x)
x = self.pool(x)
x = self.conv2(x)
x = nn.functional.relu(x)
x = self.pool(x)
x = x.view(-1, 320)
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
class CombinedNet(nn.Module):
def __init__(self):
super(CombinedNet, self).__init__()
self.net1 = Net()
self.net2 = Net()
def forward(self, x):
x1 = self.net1(x)
x2 = self.net2(x)
x = torch.cat((x1, x2), dim=1)
return x
net = CombinedNet()
input = torch.randn(1, 1, 28, 28)
output = net(input)
```
在上述代码中,定义了一个名为 `Net` 的神经网络模型,包含了两个卷积层、两个全连接层和一个池化层。然后定义了一个名为 `CombinedNet` 的神经网络模型,包含了两个 `Net` 模型,并将它们的输出在第二个维度上进行拼接。通过对 `CombinedNet` 进行前向传播,可以得到模型的输出结果。
4. CPU/GPU 转换:
```python
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(3, 10)
self.fc2 = nn.Linear(10, 1)
def forward(self, x):
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net().to(device)
input = torch.randn(1, 3).to(device)
output = net(input)
```
在上述代码中,定义了一个名为 `Net` 的神经网络模型,继承自 `nn.Module`。通过 `torch.device` 方法可以获取当前设备,然后使用 `.to(device)` 方法将模型和输入数据转换到对应的设备上。