深度学习 明星图像识别 可以套用人脸识别

############################################################# ############ 编写人：常泽钰 ############ 时间 ： 2022/9/23 ############ 功能 ：实现明星的图片分类 ############ 数据格式： ''' --dataset --train --0 --**.jpg --....... --1 --2 --3 --val ''' ############################################################# from tqdm import tqdm batch_size = 32 #设置批次大小 learning_rate = 0.001 #设置学习率 epoches = 8 #设置训练的次数 num_of_classes=12 #要分的类别个数 import torch import os from torch.utils import data import torchvision.datasets as dsets import torchvision.transforms as transforms trainpath = './dataset/train/' valpath = './dataset/val/' #数据增强的方式 traintransform = transforms .Compose([ transforms .RandomRotation (20), #随机旋转角度 transforms .ColorJitter(brightness=0.1), #颜色亮度 transforms .Resize([224, 224]), #设置成224×224大小的张量 transforms .ToTensor(), # 将图⽚数据变为tensor格式 # transforms.Normalize(mean=[0.485, 0.456, 0.406], # std=[0.229, 0.224, 0.225]), ]) valtransform = transforms .Compose([ transforms .Resize([224, 224]), transforms .ToTensor(), # 将图⽚数据变为tensor格式 ]) trainData = dsets.ImageFolder (trainpath, transform =traintransform ) # 读取训练集，标签就是train⽬录下的⽂件夹的名字，图像保存在格⼦标签下的⽂件夹⾥ valData = dsets.ImageFolder (valpath, transform =valtransform ) #读取演正剧 trainLoader = torch.utils.data.DataLoader(dataset=trainData, batch_size=batch_size, shuffle=True) #将数据集分批次 并打乱顺序 valLoader = torch.utils.data.DataLoader(dataset=valData, batch_size=batch_size, shuffle=False) #将测试集分批次并打乱顺序 # test_sum = sum([len(x) for _, _, x in os.walk(os.path.dirname(trainpath))]) #计算 训练集和测试集的图片总数 # train_sum = sum([len(x) for _, _, x in os.walk(os.path.dirname(valpath))]) train_sum=len(trainData) test_sum=len(valData) import numpy as np import torchvision.models as models # model = models.resnet34(pretrained=True) #pretrained表⽰是否加载已经与训练好的参数 # model.fc = torch.nn.Linear(512, num_of_classes) #将最后的fc层的输出改为标签数量（如3）,512取决于原始⽹络fc层的输⼊通道 model=torch.load("./resnet34_star.pt") model = model.cuda() # 如果有GPU，⽽且确认使⽤则保留；如果没有GPU，请删除 #from MyLoss import Loss #criterion=Loss("mean") criterion = torch.nn.CrossEntropyLoss() # 定义损失函数 #optimizer = torch.optim.Adam(model.parameters (), lr=learning_rate ) optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) # 定义优化器 from torch.autograd import Variable #定义训练的函数 def train(model, optimizer, criterion): model.train() total_loss = 0 train_corrects = 0 for i, (image, label) in enumerate (tqdm(trainLoader)): image = Variable(image.cuda()) # 同理 label = Variable(label.cuda()) # 同理 #print(i,image,label) optimizer.zero_grad () target = model(image) loss = criterion(target, label) loss.backward() optimizer.step() total_loss += loss.item() max_value , max_index = torch.max(target, 1) pred_label = max_index.cpu().numpy() true_label = label.cpu().numpy() train_corrects += np.sum(pred_label == true_label) return total_loss / float(len(trainLoader)), train_corrects / train_sum testLoader=valLoader #定义测试的函数 def evaluate(model, criterion): model.eval() corrects = eval_loss = 0 with torch.no_grad(): for image, label in tqdm(testLoader): image = Variable(image.cuda()) # 如果不使⽤GPU，删除.cuda() label = Variable(label.cuda()) # 同理 pred = model(image) print(pred.size()) print(label.size()) loss = criterion(pred, label) eval_loss += loss.item() max_value, max_index = torch.max(pred, 1) pred_label = max_index.cpu().numpy() true_label = label.cpu().numpy() corrects += np.sum(pred_label == true_label) return eval_loss / float(len(testLoader)), corrects, corrects / test_sum #torch.save(model,"./resnet1.pt") model.eval() corrects = eval_loss = 0 from sklearn import metrics pred_labels=np.array([]) true_labels=np.array([]) with torch.no_grad(): for image, label in tqdm(testLoader): image = Variable(image.cuda()) # 如果不使⽤GPU，删除.cuda() label = Variable(label.cuda()) # 同理 pred = model(image) print(pred.size()) print(label.size()) loss = criterion(pred, label) eval_loss += loss.item() max_value, max_index = torch.max(pred, 1) pred_label = max_index.cpu().numpy() true_label = label.cpu().numpy() corrects += np.sum(pred_label == true_label) #return eval_loss / float(len(testLoader)), corrects, corrects / test_sum pred_labels=np.append(pred_labels,pred_label) true_labels = np.append(true_labels, true_label) cm=metrics.confusion_matrix(pred_labels,true_labels,labels=[0,1,2,3,4]) print(cm) import seaborn as sns import matplotlib.pyplot as plt sns.heatmap(cm,annot=True,fmt='.2e',cmap='GnBu') plt.show() # # for i in range(epoches): # print("第{}个epoch".format(i+1)) # #train_loss,train_acc=train(model,optimizer,criterion) # #print("train_loss: {} train_acc: {}\n".format(train_loss,train_acc)) # test_loss,test_correct,test_acc=evaluate(model,criterion) # print("test_loss: {} test_correct:{} test_acc:{}".format(test_loss,test_correct,test_acc)) # torch.save(model,"./resnet34_star.pt")#保存模型，第二个参数是保存的路径