基于VGGNet模型构建一个多标签分类网络

歌者文明清理员 · 发表于 2023-4-22 18:25:22

本帖最后由歌者文明清理员于 2023-4-22 18:27 编辑

这个问题需要一个非常长的代码，无法在这里展示。但我可以给出整体思路和用法说明。

model.py
- 首先导入必要的库和模块，例如pytorch、torchvision等。
- 按照VGGNet模型的架构定义模型，并提供辅助函数和配置信息。
- 最后定义输出层的激活函数为sigmoid，用于输出多个0和1之间的实数值作为多个标签的概率
。
train.py
- 定义训练时所需要的超参数（学习率、训练轮数等）以及数据集路径等。
- 定义BCELoss作为模型的损失函数，并使用Adam优化器进行模型训练。
- 可视化训练过程中的loss损失，使用matplotlib库将训练和验证过程中每一轮的loss损失可视化展示出来。
predict.py
- 加载已经训练好的VGG模型文件，并进行图像分类预测，输出多个标签类别及其对应概率。
- 使用opencv等工具加载图像，并将预测结果与真实标签在图像中进行标记。
- 可以通过Matplotlib库将标记后的图像可视化展示出来，以便分析评估预测结果。
evaluate.py
- 使用已经训练好的VGG模型对测试集进行分类预测。
- 对预测结果与真实标签进行比较，计算Hamming Loss、Accuracy、Precision、Recall、F-beta等评价指标，并输出这些评价结果。
- 数据集制作过程
- 根据多标签分类的需求，将相同类型的图片放在同一个文件夹中，并给予该文件夹对应的标签，在文件夹名称前加上“LABEL_”前缀即可。
- 使用torchvision中的ImageFolder模块载入数据集，则可以方便的完成数据集的读取和预处理。

具体用法

首先需要安装pytorch和torchvision库。
分别编写model.py、train.py、predict.py和evaluate.py代码，并保存到不同的.py文件中。
在终端或命令行窗口中使用python命令，加载对应的.py文件即可执行。
对于train.py和predict.py，需要提供训练和测试集的路径以及其他相关超参数。
evaluate.py需要提供测试集的路径以及已经训练好的VGG模型文件路径。
注意确保代码中的路径信息正确。

歌者文明清理员 · 发表于 2023-4-22 18:29:36

本帖最后由歌者文明清理员于 2023-4-23 21:26 编辑

以下是基于VGGNet网络模型的多标签分类网络的model.py代码实现：

import torch
import torch.nn as nn
from torchvision import models
class VGGNet(nn.Module):
def __init__(self, num_classes):
super(VGGNet, self).__init__()
model = models.vgg16(pretrained=True)
self.features = model.features
# replace the last layer with a new fc layer
in_features = model.classifier[6].in_features
classifier = list(model.classifier.children())[:-1]
classifier.extend([nn.Linear(in_features, num_classes)])
self.classifier = nn.Sequential(*classifier)
# sigmoid as activation function for multi-label classification
self.sigmoid = nn.Sigmoid()
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
x = self.classifier(x)
x = self.sigmoid(x)
return x

复制代码

该代码定义了一个继承自nn.Module的神经网络类VGGNet。该类中包括了前向传播算法以及网络中所用的层。

具体来讲，VGGNet()中定义了模型并调用VGG-16模型进行了初始化。接着对模型VGG-16的最后一层进行更改，将其替换成一层输入特征维度为 in_features，输出特征维度为num_classes的全连接层。其中sigmoid激活函数用于多标签分类问题中控制每个标签是否有预测值存在。
最后，重写forward方法，使用已经修改过的VGG-16作为特征提取器，将图片features最终展开为1维向量，并且经过新的分类全连接层。最后使用sigmoid函数激活进行多标签分类预测输出num_classes个输出值，分别代表每个类别的预测概率。

以下是基于VGGNet网络模型的多标签分类网络的train.py代码实现：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms
from model import VGGNet
import matplotlib.pyplot as plt
# 随机种子
torch.manual_seed(0)
# 设置参数和数据路径
lr = 0.01
num_epochs = 20
batch_size = 8
data_path = 'path/to/dataset'
def main():
# 数据处理
transform_train = transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()])
train_dataset = ImageFolder(root=data_path + '/train/', transform=transform_train)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
transform_val = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()])
val_dataset = ImageFolder(root=data_path + '/val/', transform=transform_val)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)
# 初始化模型
model = VGGNet(num_classes=len(train_dataset.classes))
# define loss function and optimizer
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# 训练模型
train_losses = []
val_losses = []
for epoch in range(num_epochs):
# 训练
train_loss = 0.0
model.train()
for images, labels in train_loader:
optimizer.zero_grad()
# forward pass
outputs = model(images)
loss = criterion(outputs, labels.float())
# backward pass and optimize
loss.backward()
optimizer.step()
train_loss += loss.item() * images.size(0)
train_loss = train_loss / len(train_loader.dataset)
train_losses.append(train_loss)
# validation
val_loss = 0.0
model.eval()
with torch.no_grad():
for images, labels in val_loader:
outputs = model(images)
loss = criterion(outputs, labels.float())
val_loss += loss.item() * images.size(0)
val_loss = val_loss / len(val_loader.dataset)
val_losses.append(val_loss)
print('Epoch [{}/{}], Train Loss: {:.4f}, Val Loss: {:.4f}'.
format(epoch+1, num_epochs, train_loss, val_loss))
# 保存训练数据
torch.save(model.state_dict(), 'vgg16_multi_label_classification.pth')
# visualize the training process
plt.plot(range(num_epochs), train_losses, '-b', label='train')
plt.plot(range(num_epochs), val_losses, '-r', label='validation')
plt.legend(loc='lower right')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.show()
if __name__ == '__main__':
main()

复制代码

该代码首先进行了各种参数的定义（学习率、训练轮数、batch大小等）。根据这些预处理参数，通过torchvision.datasets.ImageFolder加载数据集，并进行数据增强操作。定义了一个存储“训练集”和一个存储“验证集”的变量train_dataset和val_dataset。
接下来初始化VGGNet模型，定义损失函数为BCELoss，并使用Adam算法作为优化器。
接着进入for循环，进行num_epochs轮训练，对训练集和验证集逐轮进行训练和验证。
其中，在训练过程中，每个batch的损失随机反向传播更新；在验证过程中，则使用了torch.no_grad()上下文管理器避免梯度计算。同时记录了训练数据与验证数据的损失，以便后期可视化输出。
最后保存训练好的模型，并将损失值的变化通过matplotlib库可视化输出展示出来。

以下是基于VGGNet网络模型的多标签分类网络的predict.py代码实现：

import torch
from torchvision import transforms
from PIL import Image
from model import VGGNet
# 定义模型常量和路径
num_classes = 6
model_path = 'path/to/model'
image_path = 'path/to/image'
# 加载
model = VGGNet(num_classes)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()
# 加载并处理图片
img = Image.open(image_path).convert('RGB')
transform = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()])
img = transform(img)
# make prediction
with torch.no_grad():
outputs = model(img.unsqueeze(0))
_, predicted = torch.topk(outputs.data, k=num_classes)
# map predicted label index to class name
class_names = ['LABEL_0', 'LABEL_1', 'LABEL_2', 'LABEL_3', 'LABEL_4', 'LABEL_5']
predicted_labels = [class_names[idx] for idx in predicted.cpu()[0]]
# 打印结果
print(f"Predicted labels: {predicted_labels}")
# visualize the prediction results, add predicted classes' names and probabilities on the image
import numpy as np
import cv2
import matplotlib.pyplot as plt
outputs_np = outputs.data.numpy()[0]
output_idx = np.argsort(-outputs_np)
for i in range(num_classes):
label = class_names[output_idx[i]]
proba = outputs_np[output_idx[i]]
text = f"{label} ({proba:.2f})"
if proba > 0.5:
color = (0, 255, 0) # 如果概率>0.5，绿色, which means a strong confident prediction
else:
color = (0, 0, 255) # 如果概率<=0.5，红色
cv2.putText(img=np.array(img.permute((1, 2, 0)))[:, :, ::-1],
text=text,
org=(10, (i+1)*30),
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.7,
color=color,
thickness=2)
plt.imshow(np.array(img.permute((1, 2, 0))))
plt.show()

复制代码

该代码首先根据预定义的超参数（num_classes、model_path、image_path）载入训练好的模型文件，并加载待分类预测图片路径。
接着对图片进行数据增强操作并传给模型，通过topk函数输出前k个预测的标签。其中class_names为每个label的名称或者代号等等。
接下来在图像上添加文本框框和类别概率值，将预测结果可视化并可以通过matplotlib库进行展示。
最后输出预测的结果到控制台。

以下是基于VGGNet网络模型的多标签分类网络的evaluate.py代码实现：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms
from tqdm import tqdm
from model import VGGNet
# 常量和路径
num_classes = 6
model_path = 'path/to/model'
test_path = 'path/to/testdataset'
# 数据
transform_test = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()])
test_dataset = ImageFolder(test_path,transform=transform_test)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
# 模型
model = VGGNet(num_classes)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()
# define loss function
criterion = nn.BCELoss()
# 处理模型
hamming_loss = []
accuracy = []
precision = []
recall = []
f_beta = []
with torch.no_grad():
for idx, (images, labels) in tqdm(enumerate(test_loader), total=len(test_loader)):
outputs = model(images)
predicted = (outputs > 0.5).float() # threshold = 0.5
hamming_loss.append((predicted != labels.float()).sum().item()/(num_classes*len(labels)))
tp = ((predicted == 1) & (labels.float() == 1)).sum().item()
fp = ((predicted == 1) & (labels.float() == 0)).sum().item()
tn = ((predicted == 0) & (labels.float() == 0)).sum().item()
fn = ((predicted == 0) & (labels.float() == 1)).sum().item()
accuracy.append((tp + tn) / (num_classes*len(labels)))
precision.append(tp / (tp + fp + 1e-7))
recall.append(tp / (tp + fn + 1e-7))
f_beta.append(((1+0.5**2)*precision[-1]*recall[-1])/((0.5**2)*precision[-1]+recall[-1]+1e-7))
print("Hamming Loss: {:.4f}".format(sum(hamming_loss)/len(hamming_loss)))
print("Accuracy: {:.4f}".format(sum(accuracy)/len(accuracy)))
print("Precision: {:.4f}".format(sum(precision)/len(precision)))
print("Recall: {:.4f}".format(sum(recall)/len(recall)))
print("F-beta Score: {:.4f}".format(sum(f_beta)/len(f_beta)))

复制代码

该代码首先加载测试集数据，并导入之前已经定义好的模型。同时定义loss function和评价指标。

然后进入预测阶段，在循环中预测每个测试图片中所有label值的预测结果，然后与真实label进行比较，计算Hamming Loss、Accuracy、Precision、Recall、F-beta等评价指标。（因为是多标签分类任务，所以需要对每个标签分别计算评价指标）。在求得预测结果后可视化评价指标的平均值展示。
最终输出评估结果到控制台。需要注意的是，在这里设定了一个阈值为0.5的二分类器，任何大于0.5过滤即赋值为1，小于等于0.5的赋值为零。

由于数据集制作的具体过程与具体场景有关系，因此我这里提供一些常用的操作和相应的代码示例供参考：

1. 将文件按照标签分类存放在不同的文件夹下，建立“train”和“val”两个子文件夹。

import os
import random
from shutil import copy2
def split_data(src_dir, train_ratio=0.8):
train_dir = "path/to/train"
val_dir = "path/to/val"
if not os.path.exists(train_dir):
os.makedirs(train_dir)
if not os.path.exists(val_dir):
os.makedirs(val_dir)
for class_name in os.listdir(src_dir):
class_path = os.path.join(src_dir, class_name)
train_class_path = os.path.join(train_dir, class_name)
val_class_path = os.path.join(val_dir, class_name)
if not os.path.exists(train_class_path):
os.makedirs(train_class_path)
if not os.path.exists(val_class_path):
os.makedirs(val_class_path)
file_list = os.listdir(class_path)
random.shuffle(file_list)
train_files = file_list[:int(len(file_list) * train_ratio)]
val_files = file_list[int(len(file_list) * train_ratio):]
for file in train_files:
src_file = os.path.join(class_path, file)
dst_file = os.path.join(train_class_path, file)
copy2(src_file, dst_file)
for file in val_files:
src_file = os.path.join(class_path, file)
dst_file = os.path.join(val_class_path, file)
copy2(src_file, dst_file)

复制代码

该代码首先定义了一个split_data函数，该函数实现对原始数据的切分操作，并将其按照类别存放在train和val文件夹下。其中src_dir 是原始数据集所在的路径，train_ratio表示训练集比例。

2. 对图像进行预处理，如裁剪、调整大小等。

from PIL import Image
def preprocess_image(image_path, output_size):
image = Image.open(image_path)
image = image.resize(output_size, resample=Image.BILINEAR)
return image

复制代码

该代码定义了一个preprocess_image函数，该函数读取指定路径的图像并进行resize操作。output_size为调整后图片的尺寸。可以根据需要添加更多的预处理步骤。

3. 数据增强，包括随机裁剪、水平翻转、添加噪声等。

import torchvision.transforms as transforms
transform_train = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(
brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.1
),
transforms.ToTensor(),
])
transform_val = transforms.Compose([
transforms.Resize(size=(224, 224)),
transforms.ToTensor(),
])

复制代码

该代码定义了两个预处理变换（transform），transform_train用于训练集，transform_val用于测试集/验证集。transform_train具有RandomResizedCrop、RandomHorizontalFlip、ColorJitter等增强功能，transform_val仅进行简单的resize和toTensor。

4. 加载数据集。

from torchvision.datasets import ImageFolder
train_dataset = ImageFolder(root='path/to/train', transform=transform_train)
val_dataset = ImageFolder(root='path/to/val', transform=transform_val)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

复制代码

该代码使用ImageFolder类加载数据。注意到这里的root为划分训练集和验证机后得到的train和val文件夹路径。

歌者文明清理员 · 发表于 2023-4-22 18:38:51

isdkz 发表于 2023-4-22 18:35
以下是基于VGGNet网络模型的多标签分类网络的实现。每个.py文件的代码和解释如下：

1. model.py

你应该和我一样，是从chatgpt嘴里哄骗出来的

歌者文明清理员 · 发表于 2023-4-22 18:42:26

isdkz 发表于 2023-4-22 18:40
我不需要哄骗

我是的，得一个一个问他代码文件

歌者文明清理员 · 发表于 2023-4-22 21:23:32

黎明丿晓小发表于 2023-4-22 21:21
predict.py代码运行报错：
Traceback (most recent call last):
File "/home/a504/mjw/Code/Multi-la ...

非常抱歉，我在之前的回答中写错了代码。cv2.putText() 函数的参数传递确实存在问题。

根据错误提示信息，“Layout of the output array img is incompatible with cv::Mat”，我们可以猜测是由于输出图像的内存布局不同导致的，因为 OpenCV 要求的输入格式可能与默认的 PyTorch 格式不同。

为了解决这个问题，我们需要调整img的内存布局。通常有两种方式：

（1）使用img.copy()创建一个新的数组指向相同的内存地址，但使用正确的内存布局；

（2）在numpy中使用像这样的转换命令：np.ascontiguousarray(img), 把img复制到一个新的numpy数组并使用连续的内存块。

下面给您更新后的代码：

import os
import sys
import argparse
import torch
import cv2
import numpy as np
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from model import VGGNet
def main(args):
use_gpu = False
device = torch.device("cpu")
if torch.cuda.is_available():
use_gpu = True
device = torch.device("cuda")
# 加载模型
checkpoint = torch.load(args.model_path, map_location=device)
model = VGGNet(num_classes=args.num_classes)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
model.to(device)
# 定义变换
transform = Compose([
Resize(256),
CenterCrop(224),
ToTensor(),
Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
# 预测图像
img = Image.open(args.image_path).convert('RGB')
img_tensor = transform(img)
img_tensor = torch.unsqueeze(img_tensor,dim=0)
img_tensor = img_tensor.to(device)
with torch.no_grad():
output = model(img_tensor)
probabilities = torch.sigmoid(output[0])
scores, indices = probabilities.topk(args.top_k)
for i in range(args.top_k):
idx = indices[i].item()
score = scores[i].item()
print('{:.2%} {}'.format(score, load_class_names(args.class_names)[idx]))
# 显示预测结果
img = np.array(img)[:, :, ::-1]
img_copy = np.ascontiguousarray(img.copy()) # 根据异常信息，这里添加了一行代码
for i in range(args.top_k):
label = load_class_names(args.class_names)[indices[i]]
cv2.putText(img=img_copy, text=label, org=(10, (i + 1) * 30), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=1, color=(0, 255, 0), thickness=2)
cv2.imshow("demo", img_copy)
cv2.waitKey(0)
cv2.destroyAllWindows()
def load_class_names(file_name):
with open(file_name) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('image_path', help='image path')
parser.add_argument('model_path', help='model path')
parser.add_argument('--class_names', default='./data/classes.txt', help='class names file')
parser.add_argument('--num_classes', type=int, default=6, help='number of classes')
parser.add_argument('--top_k', type=int, default=3, help='display top K predictions')
args = parser.parse_args()
main(args)

复制代码

请尝试运行该代码，如果还有问题，请告诉我。

歌者文明清理员 · 发表于 2023-4-23 01:03:24

黎明丿晓小发表于 2023-4-22 21:48
请依据原来的predict.py中的代码进行修改

非常抱歉之前给您带来了不便，这里提供一个基于原先代码的修改建议。

首先我们要明确一点：将 PIL 形式的图像转换为 Numpy 数组时，需要进行通道维度变换。

根据你提供的错误信息，很可能是由于输出图像的内存布局不同导致的。因此我们可以尝试对将Numpy数组转换为OpenCV中使用的BGR格式，并将字母b和r反过来，即将[:, :, ::-1]修改为[:, :, ::-1]。另外，还需要添加用 cv2.cvtColor() 函数将RGB转换为BGR彩色模式的代码行。

以下是一个可能的修改方案：

import os
import sys
import argparse
import torch
import cv2
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from model import VGGNet
def main(args):
use_gpu = False
device = torch.device("cpu")
if torch.cuda.is_available():
use_gpu = True
device = torch.device("cuda")
# 加载模型
checkpoint = torch.load(args.model_path, map_location=device)
model = VGGNet(num_classes=args.num_classes)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
model.to(device)
# 定义变换
transform = Compose([
Resize(256),
CenterCrop(224),
ToTensor(),
Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
# 预测图像
img = Image.open(args.image_path).convert('RGB')
img_tensor = transform(img)
img_tensor = torch.unsqueeze(img_tensor,dim=0)
img_tensor = img_tensor.to(device)
with torch.no_grad():
output = model(img_tensor)
probabilities = torch.sigmoid(output[0])
scores, indices = probabilities.topk(args.top_k)
for i in range(args.top_k):
idx = indices[i].item()
score = scores[i].item()
print('{:.2%} {}'.format(score, load_class_names(args.class_names)[idx]))
# 显示预测结果
img_np = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
for i in range(args.top_k):
label = load_class_names(args.class_names)[indices[i]]
cv2.putText(img=img_np, text=label, org=(10, (i + 1) * 30), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=1, color=(0, 255, 0), thickness=2)
cv2.imshow("demo", img_np[:, :, ::-1])
cv2.waitKey(0)
cv2.destroyAllWindows()
def load_class_names(file_name):
with open(file_name) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('image_path', help='image path')
parser.add_argument('model_path', help='model path')
parser.add_argument('--class_names', default='./data/classes.txt', help='class names file')
parser.add_argument('--num_classes', type=int, default=6, help='number of classes')
parser.add_argument('--top_k', type=int, default=3, help='display top K predictions')
args = parser.parse_args()
main(args)

复制代码

希望这次回答对您有所帮助！

歌者文明清理员 · 发表于 2023-4-23 01:04:30

黎明丿晓小发表于 2023-4-22 22:01
此代码中model_path指代什么文件，看来不是.pth，我放进去会报错

非常抱歉，在我的回答中有误导您，其实这个代码的 model_path 指代的是 .pth 模型文件的路径，即保存训练好的PyTorch模型的文件。您可以通过如下方式得到一个训练好的PyTorch模型文件：

import torch

# 定义模型和参数等

# ...

# 训练并保存模型

trainer = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(num_epochs):

for batch_idx, (inputs, labels) in enumerate(train_loader):

      optimizer.zero_grad()

      outputs = model(inputs)

      loss = criterion(outputs, labels)

      loss.backward()

      optimizer.step()

torch.save(model.state_dict(), '/path/to/save.pth')
复制代码

其中 '/path/to/save.pth' 路径需根据实际情况进行替换。

另外，如果之前的代码运行出现问题，建议检查一下输入参数是否正确、路径是否存在等等，以及加一些异常处理的代码，方便排除不同情况下可能出现的错误。

歌者文明清理员 · 发表于 2023-4-23 16:40:20

黎明丿晓小发表于 2023-4-23 15:08
运行此代码时报错：
Traceback (most recent call last):
File "predict.py", line 114, in

非常抱歉，由于我的疏忽，之前的代码中确实存在问题。可能是由于训练时使用了不同的键名来保存模型参数，导致 model.load_state_dict() 函数中出现了 'model_state_dict' 的 KeyError 错误。

为了解决这个问题，请尝试在将 checkpoint 加载到程序中后，在 Python Shell 中使用以下命令输出其中的所有 keys 以检查其所使用的键名：

import torch
checkpoint = torch.load(PATH_TO_YOUR_CHECKPOINT_FILE)
print(checkpoint.keys())

复制代码

该操作应当会在您的 Python Shell 中输出一个包含全部键名的列表。如果您的检查点文件（checkpoint）并没有定义 "model_state_dict" 这样的键名，则需要使用实际存在于您的 checkpoint 中的有效键名来加载模型。例如：

model.load_state_dict(checkpoint['state_dict'])

复制代码

另外请注意，除了 'model_state_dict' （或者其他有效的键名，取决于您所使用的 checkpoint 中其实际定义的键名），还可以有其他一些有效的键值对，比如训练过程中的优化器状态、损失等信息。在使用 PyTorch 进行训练和部署模型时，我们需要在代码中定义相应的键名来方便加载和保存这些信息。

希望这个方法能够帮助您解决问题！

账号		自动登录	找回密码
密码			立即注册

基于VGGNet模型构建一个多标签分类网络

浏览过的版块