1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38
| import torch.nn as nn
N, D_in, H, D_out = 64, 1000, 100, 10
x = torch.randn(N, D_in) y = torch.randn(N, D_out)
class TwoLayerNet(torch.nn.Module): def __init__(self, D_in, H, D_out): super(TwoLayerNet, self).__init__() self.linear1 = torch.nn.Linear(D_in, H, bias=False) self.linear2 = torch.nn.Linear(H, D_out, bias=False) def forward(self, x): y_pred = self.linear2(self.linear1(x).clamp(min=0)) return y_pred
model = TwoLayerNet(D_in, H, D_out) loss_fn = nn.MSELoss(reduction='sum') learning_rate = 1e-4 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for it in range(500): y_pred = model(x) loss = loss_fn(y_pred, y) if it%100==0 :print(it, loss.item())
optimizer.zero_grad() loss.backward() optimizer.step()
|