nn-neural network module的使用
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5)
self.conv2 = nn.Conv2d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
forward 主要定义了我们的一个计算,应该再所有子类中进行一个重写
import torch
from torch import nn
class Wangluo(nn.Module):
def __init__(self):
super().__init__()
def forward(self, input):
output = input + 1
return output
wangluo = Wangluo()
x = torch.tensor(3,0)
output = wangluo(x)
print(output)
|