https://pytorch.org/docs/stable/nn.html

import torch.nn as nn

import torch, torchvision
model = torchvision.models.resnet18(pretrained=True)
data = torch.rand(1, 3, 64, 64)
labels = torch.rand(1, 1000)
# forward pass 正向传播
prediction = model(data)
# backward pass 反向传播
loss = (prediction - labels).sum()
loss.backward()
# 优化器
optim = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
# gradient descent 梯度下降
optim.step()

创建model

class My_Model(nn.Module):
    def __init__(self, input_dim):
        super(My_Model, self).__init__()
        # 顺序容器，模块将按顺序执行
        self.layers = nn.Sequential(
            nn.Linear(input_dim, 16),
            nn.ReLU(),
            nn.Linear(16, 8),
            nn.ReLU(),
            nn.Linear(8, 1)
        )
    def forward(self, x):
        x = self.layers(x)
        x = x.squeeze(1) # (B, 1) -> (B)
        return x