报错信息:
Traceback (most recent call last):
File "D:/OpenCVLearning-main/项目实战三_停车场车位识别/main.py", line 32, in <module>
train_model()
File "D:\OpenCVLearning-main\项目实战三_停车场车位识别\train.py", line 24, in train_model
net.load_state_dict(torch.load(model_weight_path, map_location='cpu'))
File "D:\Anaconda3\envs\PyTorch\lib\site-packages\torch\serialization.py", line 699, in load
with _open_file_like(f, 'rb') as opened_file:
File "D:\Anaconda3\envs\PyTorch\lib\site-packages\torch\serialization.py", line 231, in _open_file_like
return _open_file(name_or_buffer, mode)
File "D:\Anaconda3\envs\PyTorch\lib\site-packages\torch\serialization.py", line 212, in __init__
super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: 'saved_model_weight/resnet34_pretrain_ori_low_torch_version.pth'
train.py代码
import os
import torch
import torch.nn as nn
from torchvision.models import resnet34
import torch.optim as optim
from utils.model_utils import get_dataloader, data_transform_pretrain, model_train
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def train_model():
# 获取dataloader
data_root = os.getcwd()
image_path = os.path.join(data_root, "train_data")
train_data_path = os.path.join(image_path, "train")
val_data_path = os.path.join(image_path, "test")
train_loader, validat_loader, train_num, val_num = get_dataloader(train_data_path, val_data_path,
data_transform_pretrain, batch_size=8)
# 创建模型 注意这里没指定类的个数,默认是1000类
net = resnet34()
model_weight_path = 'saved_model_weight/resnet34_pretrain_ori_low_torch_version.pth'
# 使用预训练的参数,然后进行finetune
net.load_state_dict(torch.load(model_weight_path, map_location='cpu'))
# 改变fc layer structure 把fc的输出维度改为2
in_channel = net.fc.in_features
net.fc = nn.Linear(in_channel, 2)
net.to(device)
# 模型训练配置
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0001)
epochs = 30
save_path = "saved_model_weight/resnet34_pretrain.pth"
best_acc = 0.
train_steps = len(train_loader)
model_train(net, train_loader, validat_loader, epochs, device, optimizer, loss_function, train_steps, val_num,
save_path, best_acc)
main.py代码
import torch
import os
# golb模块是用来查找符合特定规则命名的文件名的“路径+文件名”,其功能就是检索路径
import glob
import cv2
import pickle
import json
from PIL import Image
import matplotlib.pyplot as plt
from data_process import data_process
from train import train_model
from torchvision.models import resnet34
from predict_on_spot_img import predict_on_img
from predict_on_spot_video import predict_on_video
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if __name__ == '__main__':
# 如果没有停车位字典,先生成字典
if not os.path.exists('spot_dict.pickle'):
# 读入数据
test_images = [cv2.imread(path) for path in glob.glob('test_images/*.jpg')] # BGR
spot_dict = data_process(test_images[1])
else:
spot_dict = pickle.load(open('spot_dict.pickle', 'rb'))
# 读取训练好的ResNet模型
if not os.path.exists('saved_model_weight/resnet34_pretrain.pth'):
# 调用train函数训练模型
train_model()
else:
# read class_indict
json_path = './idx2class.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# 导入与训练好的模型
model = resnet34(num_classes=2).to(device)
weights_path = "saved_model_weight/resnet34_pretrain.pth"
model.load_state_dict(torch.load(weights_path))
# 预测一张停车场的图片
#img = plt.imread('test_images/scene1380.jpg')
#predict_on_img(img, spot_dict, model, class_indict)
# 预测parking_video.mp4
video_path = 'video/parking_video.mp4'
predict_on_video(video_path, spot_dict, model, class_indict)
电脑配置:2GB 的GPU,是不是要修改一些函数的参数
目前没有进行模型训练,要训练,但以现在来看是没训练成功