pytorch(二) 自定義神經網路模型
阿新 • • 發佈:2020-07-03
我們可以定義一個模型,這個模型繼承自nn.Module類。如果需要定義一個比Sequential模型更加複雜的模型,就需要定義nn.Module模型。
定義了__init__和 forward 兩個方法,就實現了自定義的網路模型。
init(),定義模型架構,實現每個層的定義。
forward(),實現前向傳播,返回y_pred
import torch class TwoLayerNet(torch.nn.Module): def __init__(self, D_in, H, D_out): """ In the constructor we instantiate two nn.Linear modules and assign them as member variables. """ super(TwoLayerNet, self).__init__() self.linear1 = torch.nn.Linear(D_in, H) self.linear2 = torch.nn.Linear(H, D_out) def forward(self, x): """ In the forward function we accept a Tensor of input data and we must return a Tensor of output data. We can use Modules defined in the constructor as well as arbitrary operators on Tensors. """ h_relu = self.linear1(x).clamp(min=0) y_pred = self.linear2(h_relu) return y_pred N, D_in, H, D_out = 64, 1000, 100, 10 x = torch.randn(N, D_in) y = torch.randn(N, D_out) model = TwoLayerNet(D_in, H, D_out) criterion = torch.nn.MSELoss(reduction='sum') optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) for t in range(500): y_pred = model(x) loss = criterion(y_pred, y) print(t, loss.item()) optimizer.zero_grad() loss.backward() optimizer.step()