PyTorch 使用心得

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