PyTorch 使用心得

模板

import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # self.hidden_layer = nn.Linear(in_dim, out_dim)

    def forward(self, x):
        # out = self.hidden_layer(x)
        return out

model = Model()
optimizer = optim.Adam(model.parameters(), lr=lr)
criticism = nn.MSELoss()

for i in range(epoch):
    optimizer.zero_grad()

    y_pred = model(x)
    loss = criticism(y_pred, y)

    loss.backward()
    optimizer.step()

# test
    

原文地址：https://www.cnblogs.com/fengyubo/p/9265844.html