基于mobilenetv3实现的分心驾驶行为检测python源码+模型(支持图片和视频检测).zip

import os import shutil import time from tqdm import tqdm import random import matplotlib.pyplot as plt import seaborn as sns import numpy as np # import pandas as pd import PIL.Image from IPython.display import Image # from sklearn.metrics import confusion_matrix import torch import torch.nn as nn import torchvision from torchvision import models, transforms, datasets from img_process import process_img path_train = "dataset/imgs/train" path_test = "dataset/imgs/test" device = torch.device("cuda:0") print(device) print(torch.cuda.get_device_name(device)) train_loader, test_loader = process_img(path_train, True) def train_model(model, criterion, optimizer, scheduler, n_epochs=5): losses = [] accuracies = [] test_accuracies = [] # set the model to train mode initially model.train() for epoch in tqdm(range(n_epochs)): since = time.time() running_loss = 0.0 running_correct = 0.0 for data in train_loader: # get the inputs and assign them to cuda inputs, labels = data inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() # forward + backward + optimize outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) loss = criterion(outputs, labels) loss.backward() optimizer.step() # calculate the loss/acc later running_loss += loss.item() running_correct += (labels == predicted).sum().item() epoch_duration = time.time() - since epoch_loss = running_loss / len(train_loader) epoch_acc = 100 / 32 * running_correct / len(train_loader) print("Epoch %s, duration: %d s, loss: %.4f, acc: %.4f" % (epoch + 1, epoch_duration, epoch_loss, epoch_acc)) losses.append(epoch_loss) accuracies.append(epoch_acc) # switch the model to eval mode to evaluate on test data model.eval() test_acc = eval_model(model) test_accuracies.append(test_acc) # re-set the model to train mode after validating model.train() scheduler.step(test_acc) since = time.time() print('Finished Training') return model, losses, accuracies, test_accuracies def eval_model(model): correct = 0.0 total = 0.0 with torch.no_grad(): for i, data in enumerate(test_loader, 0): images, labels = data images = images.to(device) labels = labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() test_acc = 100.0 * correct / total print('Accuracy of the network on the test images: %d %%' % ( test_acc)) return test_acc if __name__=='__main__': model_ft = models.mobilenet_v3_small(pretrained=True) num_ftrs = model_ft.classifier[3].in_features model_ft.classifier[3] = nn.Linear(num_ftrs, 10) # No. of classes = 10 model_ft = model_ft.to(device) criterion = nn.CrossEntropyLoss().to(device) optimizer = torch.optim.SGD(model_ft.parameters(), lr=0.01, momentum=0.9) lrscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=3, threshold=0.9) model_ft, training_losses, training_accs, test_accs = train_model(model_ft, criterion, optimizer, lrscheduler, n_epochs=5) torch.save(model_ft.state_dict(), "checkpoint/model-driver")