1.dataset代码 import random from torchvision import datasets,transforms from torchvision.transforms import functional_pil as F_pil from torch.utils.data import DataLoader import torch.nn as nn from PIL import Image from torchvision.transforms import functional as F IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp') class MergeTwoImageTransformer(nn.Module): def __init__(self,image_paths): super(MergeTwoImageTransformer, self).__init__() self.image_paths=image_paths def forward(self,img): if not F_pil._is_pil_image(img): raise ValueError('图像合并暂时仅支持pillow对象') img_w,img_h=F._get_image_size(img) #随机选择一个路径 img_path=random.choice(self.image_paths) #图像加载 other_img=datasets.folder.pil_loader(img_path) #图像剪切 i,j,h,w=transforms.RandomResizedCrop.get_params(other_img,scale=(0.5,1.0),ratio=(3./4.,4./3.)) other_img=F.resized_crop(other_img,i,j,h,w,(img_h,img_w),transforms.InterpolationMode.BILINEAR) # 图像合并 img=Image.blend(img,other_img,0.15) return img class Dataset: def __init__(self,root_dir,batch_size=8,num_works=0,train=True,shuffle=None): if shuffle is None: shuffle = train self.root_dir=root_dir _, class_to_idx = datasets.folder.find_classes(self.root_dir) image_paths = datasets.ImageFolder.make_dataset( self.root_dir, class_to_idx, IMG_EXTENSIONS ) self.image_paths = [s[0] for s in image_paths] transform=self.get_train_transform() if train else self.get_valid_transform() self.dataset=datasets.ImageFolder( root=self.root_dir, transform=transform ) self.loader=DataLoader( dataset=self.dataset, shuffle=shuffle, num_workers=num_works, batch_size=batch_size, prefetch_factor=2 if num_works==0 else batch_size*num_works) def __len__(self): return len(self.dataset.imgs) def __iter__(self): for data in self.loader: yield data def get_train_transform(self): return transforms.Compose([ MergeTwoImageTransformer(self.image_paths), transforms.RandomHorizontalFlip(p=0.4), transforms.RandomResizedCrop(size=(224,224),scale=(0.6,1.0)), transforms.ColorJitter(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) @staticmethod def get_valid_transform(): return transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])
2.测试准确率代码
import torch import torch.nn as nn class AccuracyScore(nn.Module): def __init__(self): super(AccuracyScore, self).__init__() # noinspection PyMethodMayBeStatic def forward(self, y_pred, y_true): y_pred_dim = y_pred.dim() y_true_dim = y_true.dim() if y_pred_dim == y_true_dim: pass elif y_pred_dim == y_true_dim + 1: y_pred = torch.argmax(y_pred, dim=1) else: raise ValueError("格式异常!") y_pred = y_pred.to(y_true.dtype) correct = (y_pred == y_true) return torch.mean(correct.to(torch.float32))
3.模型训练测试代码
import torch.nn as nn from torchvision import models import torch from image_classify_model import dataset from torch.utils.tensorboard import SummaryWriter import os from image_classify_model.metric import AccuracyScore import torch.optim as optim import numpy as np from . import network # os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID' # os.environ['CUDA_VISIBLE_DEVICES']='1' class Network(nn.Module): def __init__(self,num_classes=2,**kwargs): super(Network, self).__init__() vgg16=models.vgg16_bn(pretrained=True) for module in vgg16.parameters(): module.requires_grad=False vgg16.classifier[3]=nn.Linear(4096,64) vgg16.classifier[6]=nn.Linear(64,num_classes) self.vgg16=nn.DataParallel(vgg16) def forward(self,x): return self.vgg16(x) class ImageClassifiModel(object): def __init__(self,num_classes=2,summary_dir=None,model_dir=None,lr=0.005,momentum=0.5,weight_decay=0.05,nesterov=True,batch_size=16,num_works=0,use_gpu=True): super(ImageClassifiModel, self).__init__() self.device = torch.device('cuda' if use_gpu and torch.cuda.is_available() else 'cpu') self.net=Network(num_classes=num_classes) self.net=self.net.to(self.device) self.loss_fn=nn.CrossEntropyLoss() self.loss_fn.to(self.device) self.metrics={ 'acc':AccuracyScore() } self.train_optim=optim.SGD(params=[p for p in self.net.parameters() if p.requires_grad], lr=lr,momentum=momentum,nesterov=nesterov,weight_decay=weight_decay) self.batch_size=batch_size self.num_works=num_works if not os.path.exists(model_dir): os.makedirs(model_dir) self.model_dir = model_dir if not os.path.exists(summary_dir): os.makedirs(summary_dir) writer=SummaryWriter(log_dir=summary_dir) writer.add_graph(self.net,torch.empty(self.batch_size,3,224,224)) writer.close() self.summary_dir=summary_dir def training(self,train_data_dir,test_data_dir,total_epoch,summary_step_interval=200,eval_epoch_interval=1,save_epoch_interval=10): trainset=dataset.Dataset(root_dir=train_data_dir,batch_size=self.batch_size,num_works=self.num_works,train=True,shuffle=True) testset = dataset.Dataset(root_dir=test_data_dir, batch_size=self.batch_size, num_works=self.num_works, train=False, shuffle=False) writer = SummaryWriter(log_dir=os.path.join(self.summary_dir,'training')) total_samples = len(trainset) train_step = 0 test_step = 0 for epoch in range(total_epoch): # 训练操作 self.net.train(True) train_loss = [] for data in trainset: inputs, labels = data inputs, labels = inputs.to(self.device), labels.to(self.device) # 前向过程 outputs = self.net(inputs) _loss = self.loss_fn(outputs, labels) _metrics={} for _key in self.metrics: _metrics[_key]=self.metrics[_key](outputs,labels).cpu().numpy() self.train_optim.zero_grad() _loss.backward() self.train_optim.step() train_loss.append(_loss.item()) if train_step % summary_step_interval == 0: # 可视化输出 writer.add_scalar('train_loss', _loss, train_step) writer.add_scalars('train_metrics', _metrics, train_step) print(f"Train {epoch + 1}/{total_epoch} {train_step} " f"loss:{_loss.item():.3f} accuracy:{_metrics.get('acc',-0.0):.3f}") train_step += 1 # 测试操作 if epoch%eval_epoch_interval==0: self.net.eval() test_loss = [] for data in testset: inputs, labels = data inputs,labels=inputs.to(self.device),labels.to(self.device) # 前向过程 outputs = self.net(inputs) _loss = self.loss_fn(outputs, labels) _metrics = {} for _key in self.metrics: _metrics[_key] = self.metrics[_key](outputs, labels).cpu().numpy() test_loss.append(_loss.item()) if test_step % summary_step_interval == 0: # 可视化输出 writer.add_scalar('test_loss', _loss, test_step) writer.add_scalars('test_metrics', _metrics, test_step) print(f"Test {epoch + 1}/{total_epoch} {test_step} " f"loss:{_loss.item():.3f} accuracy:{_metrics.get('acc',-0.0):.3f}") test_step += 1 # 每个epoch计算损失 writer.add_scalars('epoch_loss', {'train': np.mean(train_loss), 'test': np.mean(test_loss)}, epoch) else: writer.add_scalars('epoch_loss', {'train': np.mean(train_loss)}, epoch) if epoch%save_epoch_interval==0: # self.net.remove_hook_fn() torch.save(self.net,os.path.join(self.model_dir,f'model{epoch}.pt')) # self.net.add_hook_fn() torch.save(self.net,os.path.join(self.model_dir,f'model{total_epoch}.pt')) writer.close() def eval(self): pass
4.main函数
root_dir=r'D:\pythonProject\0807' m=model.ImageClassifiModel(model_dir=os.path.join(root_dir,'output01','model'), summary_dir=os.path.join(root_dir,'output01','summary'), batch_size=16, use_gpu=True) root_dir=r'D:\pythonProject\dogcat' m.training(save_epoch_interval=2, total_epoch=10, eval_epoch_interval=2, summary_step_interval=100, train_data_dir=os.path.join(root_dir,'train'), test_data_dir=os.path.join(root_dir,'test') )