首先要有一個pytorch模型，我這里選用googelnet 為例，

我們可以使用pytorch 提供的imagenet的預訓練模型。

import torchvision googlenet = torchvision.models.googlenet(pretrained=True) input = torch.randn(2,3,224,224) out = googlenet(input)#控制臺會下載預訓練模型，找到下載的模型，然后直接使用

如果想自己訓練的話，就接著往下看，否則閱讀結束。

我是用的羅浩提供的框架，訓練了一個貓狗分類，下面是羅浩的框架的github，以及B 站教學視頻

https://github.com/michuanhaohao/deep-person-reid https://www.bilibili.com/video/BV1Pg4y1q7sN

第一步我們要準備貓狗數據集

鏈接：https://pan.baidu.com/s/1HBewIgKsFD8hh3ICOnnTwA 提取碼：ktab 順便直接下載我的源代碼吧。 鏈接: https://pan.baidu.com/s/1l6mrSpbfNSOsbmw2FT0zYw 提取碼: r799 來自 https://blog.csdn.net/qq_43391414/article/details/118462136

然后就是改羅浩老師的框架了

第一步先寫googlenet 網絡，在models文件夾新建GoogLeNet.py

from __future__ import absolute_importimport torch from torch import nn from torch.nn import functional as F import torchvision __all__ = ['GoogLeNet']class GoogLeNet(nn.Module):def __init__(self, num_classes=2, loss={'xent'}, **kwargs):super(GoogLeNet, self).__init__()self.loss = lossgooglenet = torchvision.models.googlenet(pretrained=True)# self.base = googlenetself.base = nn.Sequential(*list(googlenet.children())[:-2])self.classifier = nn.Linear(1024, num_classes)self.feat_dim = 1024 # feature dimensiondef forward(self, x):x = self.base(x)# x = F.avg_pool2d(x, x.size()[2:])f = x.view(x.size(0), -1)y = self.classifier(f)return yif not self.training:return fy = self.classifier(f)if self.loss == {'xent'}:return yelif self.loss == {'xent', 'htri'}:return y, felif self.loss == {'cent'}:return y, felif self.loss == {'ring'}:return y, felse:raise KeyError("Unsupported loss: {}".format(self.loss))if __name__ == "__main__":input = torch.randn(2,3,224,224)model = GoogLeNet()out = model(input)aaa= 100

在models/__init__.py修改相應配置

最上面加上 from .GoogLeNet import *factory 加上 'GoogLeNet':GoogLeNet,

第二步修改數據集相關代碼

在data_manage.py新增class

class CatDog(object):dataset_dir = 'dog_cat_dataset'def __init__(self, root='E:\\workspace\\dataset', **kwargs):self.dataset_dir = osp.join(root, self.dataset_dir)self.train_dir = osp.join(self.dataset_dir, 'train')self.test_dir = osp.join(self.dataset_dir, 'test')self.class_num = 2self._check_before_run()train, num_train_imgs = self._process_dir(self.train_dir)test, num_test_imgs = self._process_dir(self.test_dir)num_total_imgs = num_train_imgs + num_test_imgsprint("=> Dog Cat dataset loaded")print("Dataset statistics:")print(" ------------------------------")print(" subset | # images")print(" ------------------------------")print(" train | {:8d}".format(num_train_imgs))print(" test | {:8d}".format(num_test_imgs))print(" ------------------------------")print(" total | {:8d}".format(num_total_imgs))print(" ------------------------------")self.train = trainself.test = testdef _check_before_run(self):"""Check if all files are available before going deeper"""if not osp.exists(self.dataset_dir):raise RuntimeError("'{}' is not available".format(self.dataset_dir))if not osp.exists(self.train_dir):raise RuntimeError("'{}' is not available".format(self.train_dir))if not osp.exists(self.test_dir):raise RuntimeError("'{}' is not available".format(self.test_dir))def _process_dir(self, dir_path):img_paths = glob.glob(osp.join(dir_path, '*.jpg'))dataset = []for img_path in img_paths:img_path = img_path.replace('\\', '/')class_name = img_path.split(".")[0].split("/")[-1]if not class_name in ["dog","cat"]: continueif class_name == "dog":class_index = 0elif class_name == "cat":class_index = 1dataset.append((img_path, class_index))num_imgs = len(dataset)return dataset, num_imgs

修改一下factory,新增條目

__img_factory = {'market1501': Market1501,'cuhk03': CUHK03,'dukemtmcreid': DukeMTMCreID,'msmt17': MSMT17,'cat_dog':CatDog, }

一些解釋：

?root是dataset_dir 的路徑，上面是數據集的文件夾，看下的我的文件結構，訓練集和測試集隨便分分就好，如果用我的代碼，文件夾的名稱和我一致。

?在dataset_loader.py新增代碼

class DogCatDataset_test(Dataset):"""Image Person ReID Dataset"""def __init__(self, dataset, transform=None):self.dataset = datasetself.transform = transformdef __len__(self):return len(self.dataset)def __getitem__(self, index):img_path, class_index = self.dataset[index]img = read_image(img_path)if self.transform is not None:img = self.transform(img)return img_path.split("/")[-1],img, class_index class DogCatDataset(Dataset):"""Image Person ReID Dataset"""def __init__(self, dataset, transform=None):self.dataset = datasetself.transform = transformdef __len__(self):return len(self.dataset)def __getitem__(self, index):img_path, class_index = self.dataset[index]img = read_image(img_path)if self.transform is not None:img = self.transform(img)return img, class_index

訓練和測試腳本：

from __future__ import print_function, absolute_import import os import sys import time import datetime import argparse import os.path as osp import numpy as npimport torch import torch.nn as nn import torch.backends.cudnn as cudnn from torch.utils.data import DataLoader from torch.optim import lr_schedulerimport data_manager from dataset_loader import ImageDataset, DogCatDataset, DogCatDataset_test import transforms as T import models from losses import CrossEntropyLabelSmooth, DeepSupervision, CrossEntropy_loss from utils import AverageMeter, Logger, save_checkpoint from eval_metrics import evaluate from optimizers import init_optimparser = argparse.ArgumentParser(description='Train image model with cross entropy loss') # Datasets parser.add_argument('--root', type=str, default='E:\\workspace\\dataset', help="root path to data directory") parser.add_argument('-d', '--dataset', type=str, default='cat_dog',choices=data_manager.get_names()) parser.add_argument('-j', '--workers', default=4, type=int,help="number of data loading workers (default: 4)") parser.add_argument('--height', type=int, default=224,help="height of an image (default: 256)") parser.add_argument('--width', type=int, default=224,help="width of an image (default: 128)") parser.add_argument('--split-id', type=int, default=0, help="split index") # CUHK03-specific setting parser.add_argument('--cuhk03-labeled', action='store_true',help="whether to use labeled images, if false, detected images are used (default: False)") parser.add_argument('--cuhk03-classic-split', action='store_true',help="whether to use classic split by Li et al. CVPR'14 (default: False)") parser.add_argument('--use-metric-cuhk03', action='store_true',help="whether to use cuhk03-metric (default: False)") # Optimization options parser.add_argument('--optim', type=str, default='adam', help="optimization algorithm (see optimizers.py)") parser.add_argument('--max-epoch', default=60, type=int,help="maximum epochs to run") parser.add_argument('--start-epoch', default=0, type=int,help="manual epoch number (useful on restarts)") parser.add_argument('--train-batch', default=128, type=int,help="train batch size") parser.add_argument('--test-batch', default=1, type=int, help="test batch size") parser.add_argument('--lr', '--learning-rate', default=0.0003, type=float,help="initial learning rate") parser.add_argument('--stepsize', default=20, type=int,help="stepsize to decay learning rate (>0 means this is enabled)") parser.add_argument('--gamma', default=0.1, type=float,help="learning rate decay") parser.add_argument('--weight-decay', default=5e-04, type=float,help="weight decay (default: 5e-04)") # Architecture parser.add_argument('-a', '--arch', type=str,default='GoogLeNet',# default='resnet50',choices=models.get_names()) # Miscs parser.add_argument('--print-freq', type=int, default=10, help="print frequency") parser.add_argument('--seed', type=int, default=1, help="manual seed") parser.add_argument('--resume', type=str,# default='E:/workspace/classify/checkpoint_ep60.pth',metavar='PATH') parser.add_argument('--evaluate',# default=1,action='store_true', help="evaluation only") parser.add_argument('--eval-step', type=int, default=-1,help="run evaluation for every N epochs (set to -1 to test after training)") parser.add_argument('--start-eval', type=int, default=0, help="start to evaluate after specific epoch") parser.add_argument('--save-dir', type=str, default='log_resnet_dog') parser.add_argument('--use-cpu', action='store_true', help="use cpu") parser.add_argument('--gpu-devices', default='0', type=str, help='gpu device ids for CUDA_VISIBLE_DEVICES')args = parser.parse_args()def main():torch.manual_seed(args.seed)os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devicesuse_gpu = torch.cuda.is_available()if args.use_cpu: use_gpu = Falseif not args.evaluate:sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))else:sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))print("==========\nArgs:{}\n==========".format(args))if use_gpu:print("Currently using GPU {}".format(args.gpu_devices))cudnn.benchmark = Truetorch.cuda.manual_seed_all(args.seed)else:print("Currently using CPU (GPU is highly recommended)")print("Initializing dataset {}".format(args.dataset))dataset = data_manager.init_img_dataset(root=args.root, name=args.dataset, split_id=args.split_id,cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split,)transform_train = T.Compose([T.Random2DTranslation(args.height, args.width),T.RandomHorizontalFlip(),T.ToTensor(),# T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])transform_test = T.Compose([T.Resize((args.height, args.width)),T.ToTensor(),# T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])pin_memory = True if use_gpu else Falsetrainloader = DataLoader(DogCatDataset(dataset.train, transform=transform_train),batch_size=args.train_batch, shuffle=True, num_workers=args.workers,pin_memory=pin_memory, drop_last=True,)testloader = DataLoader(DogCatDataset_test(dataset.test, transform=transform_test),batch_size=args.test_batch, shuffle=False, num_workers=args.workers,pin_memory=pin_memory, drop_last=False,)print("Initializing model: {}".format(args.arch))model = models.init_model(name=args.arch, num_classes=2, loss={'xent'}, use_gpu=use_gpu)print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters()) / 1000000.0))# criterion = CrossEntropyLabelSmooth(num_classes=dataset.class_num, use_gpu=use_gpu)criterion = CrossEntropy_loss(num_classes=dataset.class_num, use_gpu=use_gpu)optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)if args.stepsize > 0:scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)start_epoch = args.start_epochif args.resume:print("Loading checkpoint from '{}'".format(args.resume))checkpoint = torch.load(args.resume,map_location=torch.device('cpu'))model.load_state_dict(checkpoint)# start_epoch = checkpoint['epoch']if use_gpu:model = nn.DataParallel(model).cuda()if args.evaluate:print("Evaluate only")test(model, testloader, use_gpu)returnstart_time = time.time()train_time = 0best_rank1 = -np.infbest_epoch = 0print("==> Start training")for epoch in range(start_epoch, args.max_epoch):start_train_time = time.time()train(epoch, model, criterion, optimizer, trainloader, use_gpu)train_time += round(time.time() - start_train_time)if args.stepsize > 0: scheduler.step()if use_gpu:state_dict = model.module.state_dict()else:state_dict = model.state_dict()save_checkpoint(state_dict, 0, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth'))# if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (# epoch + 1) == args.max_epoch:# print("==> Test")# # rank1 = test(model, testloader, use_gpu)# # is_best = rank1 > best_rank1# # if is_best:# # best_rank1 = rank1# # best_epoch = epoch + 1## if use_gpu:# state_dict = model.module.state_dict()# else:# state_dict = model.state_dict()## save_checkpoint(state_dict, 0, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))elapsed = round(time.time() - start_time)elapsed = str(datetime.timedelta(seconds=elapsed))train_time = str(datetime.timedelta(seconds=train_time))print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))def train(epoch, model, criterion, optimizer, trainloader, use_gpu):losses = AverageMeter()batch_time = AverageMeter()data_time = AverageMeter()model.train()end = time.time()for batch_idx, (imgs, class_index, ) in enumerate(trainloader):if use_gpu:imgs, class_index = imgs.cuda(), class_index.cuda()# measure data loading timedata_time.update(time.time() - end)outputs = model(imgs)if isinstance(outputs, tuple):loss = DeepSupervision(criterion, outputs, class_index)else:loss = criterion(outputs, class_index)optimizer.zero_grad()loss.backward()optimizer.step()# measure elapsed timebatch_time.update(time.time() - end)end = time.time()losses.update(loss.item(), class_index.size(0))if (batch_idx + 1) % args.print_freq == 0:print('Epoch: [{0}][{1}/{2}]\t''Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t''Data {data_time.val:.3f} ({data_time.avg:.3f})\t''Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch + 1, batch_idx + 1, len(trainloader), batch_time=batch_time,data_time=data_time, loss=losses))def test(model, testloader,use_gpu):batch_time = AverageMeter()model.eval()with torch.no_grad():for batch_idx, (file_name,imgs, class_index) in enumerate(testloader):if use_gpu: imgs = imgs.cuda()end = time.time()result = model(imgs)batch_time.update(time.time() - end)print(file_name,result ," true",class_index)np.save(file_name[0],imgs.numpy())return 0if __name__ == '__main__':main()

parser.add_argument('--root', type=str, default='E:\\workspace\\dataset', help="root path to data directory")這里改成自己的文件夾

拿到訓練結果一個pth

如何測試pth呢

?只需改

然后重新運行腳本

?

?前面是文件名，后面是推理結果score,? 貓的index是1 （如最后一列所示），可以看到后面那列的數據更大，說明預測的對

我的的代碼地址 https://gitee.com/feboreigns/classify

總結

以上是生活随笔為你收集整理的教你自己训练的pytorch模型转caffe(一）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。