一、搭建模型的步骤

在 PyTorch 中，可以使用 torch.nn 模块来搭建深度学习模型。具体步骤如下：

1. 定义一个继承自 torch.nn.Module 的类，这个类将作为我们自己定义的模型。

2. 在类的构造函数 __init__() 中定义网络的各个层和参数。可以使用 torch.nn 模块中的各种层，如 Conv2d、BatchNorm2d、Linear 等。

3. 在类中定义前向传播函数 forward()，实现模型的具体计算过程。

4. 将模型部署到 GPU 上，可以使用 model.to(device) 将模型移动到指定的 GPU 设备上。

二、简单的例子

下面是一个简单的例子，演示了如何使用 torch.nn 模块搭建一个简单的全连接神经网络：

import torch.nn as nn

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

MyNet 的神经网络类，它继承自 torch.nn.Module。在构造函数 __init__() 中定义了两个全连接层，一个 ReLU 激活函数，并将它们作为网络的成员变量。在前向传播函数 forward() 中，首先将输入的图像数据 x 压成一维向量，然后依次经过两个全连接层和一个 ReLU 激活函数，最终得到模型的输出结果。

在模型训练之前，需要将模型部署到 GPU 上，可以使用以下代码将模型移动到 GPU 上：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = MyNet().to(device)

如何将loss函数添加到模型中去呢？

在 PyTorch 中，通常将损失函数作为单独的对象来定义，并在训练过程中手动计算和优化损失。为了将损失函数添加到模型中，需要在模型类中添加一个成员变量，然后在前向传播函数中计算损失。

下面是一个例子，演示了如何在模型中添加交叉熵损失函数：

import torch.nn as nn

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(512, 10)
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, x, y):
        x = x.view(-1, 784)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        loss = self.loss_fn(x, y)
        return x, loss

在模型类 MyNet 的构造函数中添加了一个成员变量 self.loss_fn，它是交叉熵损失函数。在前向传播函数 forward() 中，传入两个参数 x 和 y，其中 x 是输入图像数据，y 是对应的标签。在函数中先执行正向传播计算，然后计算交叉熵损失，并将损失值作为输出返回。

实际训练代码

在实际训练过程中，首先将模型输出结果 x 和标签 y 传入前向传播函数 forward() 中计算损失，然后使用优化器更新模型的权重和偏置。代码如下：

model = MyNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for inputs, labels in data_loader:
    inputs = inputs.to(device)
    labels = labels.to(device)
    optimizer.zero_grad()
    outputs, loss = model(inputs, labels)
    loss.backward()
    optimizer.step()

在上面的代码中，使用随机梯度下降优化器 torch.optim.SGD 来更新模型的参数。在每个批次中，首先将输入数据和标签移动到 GPU 上，然后使用 optimizer.zero_grad() 将梯度清零。接着执行前向传播计算，并得到损失值 loss。最后使用 loss.backward() 计算梯度并执行反向传播，使用 optimizer.step() 更新模型参数。

2023.03.27更新 完整的代码

# -*-coding:utf-8-*-
# !/usr/bin/env python
# @Time    : 2023/3/27 上午11:00
# @Author  : loveinfall uestc
# @File    : csdn_test_.py
# @Description :

import torch
import torch.nn as nn
import torch.utils.data as data
import cv2

####################### model ###########################
class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x
###################### end ##############################

################# loss 函数 #############################
loss_fn = nn.CrossEntropyLoss()
################## end ################################


#################### dataloader 需要自己构建 ############
class image_folder(data.Dataset):
    def __init__(self):
        self.image_dirs = []#构造数据读取路径列表
        self.label_dirs = []
    def __getitem__(self,index):
        image = cv2.imread(self.image_dirs[index])
        label = 'read data'#根据实际情况，写
        return  image,label
    def __len__(self):
        return 'len(data)'
    
train_dataset = image_folder()

data_loader = data.DataLoader(
        train_dataset,
        batch_size=3,
        shuffle=True,
        num_workers=2,
        pin_memory=True)
#################### end ################################


##################### train #######################@#####
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = MyNet().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for inputs, labels in data_loader:
    inputs = inputs.to(device)
    labels = labels.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs,labels)
    loss.backward()
    optimizer.step()