快速上手
用熟悉的 Module、优化器与反向传播接口,完成一次训练循环。
定义、训练、更新
import tensorplay as tp
import tensorplay.nn as nn
import tensorplay.optim as optim
class Net(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 1)
def forward(self, x):
return self.fc2(tp.relu(self.fc1(x)))
model = Net()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
inputs = tp.randn(32, 10)
target = tp.randn(32, 1)
for epoch in range(100):
optimizer.zero_grad()
output = model(inputs)
loss = criterion(output, target)
loss.backward()
optimizer.step()