總結一下最近整合的兩段訓練的程式碼(貓狗圖片分類)以供以後學習參考使用
阿新 • • 發佈：2021-11-02
第一種程式碼
點選檢視程式碼
import os import random import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader,Dataset import torchvision.transforms as transforms import torch.optim as optim from PIL import Image def set_seed(seed=1): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) set_seed() # 設定隨機種子-使結果具有重複性，重現結果，這功能就是為了DataLoader中的shuffle，使每一次訓練都亂的一樣 #功能上，torch.manual_seed(0)=set_seed() pet_label={'cat':0,'dog':1} #引數設定 Max_epoch=5 Batch_size=100 Lr=0.02 #建立MyDataset類 class MyDataset(Dataset): def __init__(self,data_dir,transform=None): self.label_name={'cat':0,'dog':1} self.data_info = self.get_img_info(data_dir) self.transform = transform def __getitem__(self, index): path_img, label = self.data_info[index] img = Image.open(path_img).convert('RGB') if self.transform is not None: img = self.transform(img) return img, label def __len__(self): return len(self.data_info) @staticmethod #靜態方法繫結到一個類而不是該類物件，意味著可以在沒有該類物件的情況下呼叫靜態方法。 def get_img_info(data_dir): data_info=list() for root,dirs,_ in os.walk(data_dir): for sub_dir in dirs: img_names=os.listdir(os.path.join(root,sub_dir)) for i in range(len(img_names)): img_name=img_names[i] path_img=os.path.join(root,sub_dir,img_name) label=pet_label[sub_dir] data_info.append((path_img,int(label))) return data_info #定義網路 class Net(nn.Module): def __init__(self,classes): super(Net,self).__init__() self.conv1=nn.Sequential( nn.Conv2d(3,6,5), nn.ReLU(), nn.MaxPool2d(2,2) ) self.conv2=nn.Sequential( nn.Conv2d(6,16,5), nn.ReLU(), nn.MaxPool2d(2,2) ) self.fc=nn.Sequential( nn.Linear(16*5*5,120), nn.ReLU(), nn.Linear(120,84), nn.ReLU(), nn.Linear(84,classes) ) def forward(self,x): out=self.conv1(x) out=self.conv2(out) out=out.view(out.size(0),-1) out=self.fc(out) return out #------------------------------------step 1 資料-------------------------------------------- train_dir='/home/yanhua/Documents/catdogxunlian/dogcat/train' test_dir='/home/yanhua/Documents/catdogxunlian/dogcat/test' norm_mean = [0.485, 0.456, 0.406] norm_std = [0.229, 0.224, 0.225] common_transform=transforms.Compose([ transforms.Resize((32,32)), transforms.ToTensor(), transforms.Normalize(norm_mean,norm_std) ]) #構建MyDataset例項 train_data=MyDataset(data_dir=train_dir,transform=common_transform) test_data=MyDataset(data_dir=test_dir,transform=common_transform) #構建DataLoder train_loader=DataLoader(dataset=train_data,batch_size=Batch_size,shuffle=True) test_loader=DataLoader(dataset=test_data,batch_size=Batch_size) #------------------------------------step 2 模型------------------------------------------- net=Net(classes=2) #------------------------------------step 3 損失函式--------------------------------------- criterion=nn.CrossEntropyLoss() #------------------------------------step 4 優化器----------------------------------------- optimizer=optim.SGD(net.parameters(),lr=Lr,momentum=0.8) #------------------------------------step 5 訓練------------------------------------------- for epoch in range(Max_epoch): loss_mean=0. correct=0. total=0. for step,(batch_x,batch_y) in enumerate(train_loader): #因為enumerate的返回特性，才有step這個變數 optimizer.zero_grad() outputs=net(batch_x) loss=criterion(outputs,batch_y) loss.backward() optimizer.step() _, predicted = torch.max(outputs.data, 1) total+=batch_y.size(0) correct+=(predicted ==batch_y).squeeze().sum().numpy() loss_mean+=loss.item() if (step+1) %5 == 0: loss_mean = loss_mean / 5 print("Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(epoch, Max_epoch, step+1, len(train_loader), loss_mean, correct / total)) loss_mean = 0. #這裡correct和total並不需要清零，因為這兩哥們都是共進退的 #torch.save(net.state_dict(),'/home/yanhua/Documents/catdogxunlian/lenet.pth')--儲存模型引數
第二種程式碼
點選檢視程式碼
from torchvision.transforms import transforms from torchvision.datasets import ImageFolder #ImageFolder用於載入圖片與相應的標籤 import torch from torch.utils.data import DataLoader from torchvision import models import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler #1、準備資料 simple_transform=transforms.Compose( [ transforms.Resize((224,224)), #這裡一定要注意((224,224))的兩個括號或([]) transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]) ] ) train=ImageFolder('dogsandcats/train/',simple_transform) #train和valid兩個物件包含了類別和相應資料集索引的對映 valid=ImageFolder('dogsandcats/valid/',simple_transform) #print(train.class_to_idx) #print(train.classes) train_data_gen=DataLoader(train,batch_size=100,num_workers=3,shuffle=True)#num_workers負責併發 valid_data_gen=DataLoader(valid,batch_size=100,num_workers=3) dataloaders={'train':train_data_gen,'valid':valid_data_gen} dataset_sizes={'train':len(train_data_gen),'valid':len(valid_data_gen)} torch.manual_seed(0) #2、構建網路 model_ft=models.resnet18(pretrained=True) #pretrained是否使用預訓練好的權重(ImageNet分類問題) num_ftrs=model_ft.fc.in_features model_ft.fc=nn.Linear(num_ftrs,2) if torch.cuda.is_available(): model_ft=model_ft.cuda() #3、選擇loss和優化器 learning_rate=0.001 criterion=nn.CrossEntropyLoss() optimizer_ft=optim.SGD(model_ft.parameters(),lr=learning_rate,momentum=0.9) exp_lr_scheduler=lr_scheduler.StepLR(optimizer_ft,step_size=7,gamma=0.1) #StepLR幫助動態修改學習率 #4、訓練模型 def train_model(model,criterion,optimizer,scheduler,num_epochs): for epoch in range(num_epochs): #每輪帶有訓練和驗證階段 for phase in ['train','valid']: if phase == 'train': model.train(True) # 模型設為訓練模式——default=True，啟用Batch Normalization和Dropout else: model.train(False) # 模型設為評估模式model.train(False)=model.eval(),不啟用BN和Dropout running_loss=0.0 running_corrects=0.0 total=0 #在資料上迭代 for step,(inputs,labels) in enumerate(dataloaders[phase]): if torch.cuda.is_available(): inputs,labels=inputs.cuda(),labels.cuda() #梯度引數清零 optimizer.zero_grad() #前向傳播 outputs=model(inputs) loss=criterion(outputs,labels) #只在訓練階段反向優化 if phase =='train': loss.backward() optimizer.step() scheduler.step() #統計 _,preds=torch.max(outputs.data,1) ##torch.max(tensor,dim)-dim:按索引消去維度,不加_,返回的是一行中最大的數,加_,則返回一行中最大數的位置 running_loss+=loss.item() #running_corrects+=torch.sum(preds==labels.data) running_corrects += (preds == labels).squeeze().sum().numpy() total+=labels.size(0) if (step+1)%5==0: loss_mean=running_loss/5 print("{}:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(phase,epoch,num_epochs,step + 1,len(dataloaders[phase]),loss_mean,running_corrects/ total)) loss_mean=0.0 running_loss=0.0 if __name__=='__main__': train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=10)
而有shuffle時，這Acc是真的穩啊，資料也很好看，能清楚地發現Acc在穩步上升。
程式碼可能還有不完善的地方，不過先到這裡吧。