Pytorch图像标签分类之EfficientNet-PyTorch

GTX_AI

Pytorch图像标签分类之EfficientNet-PyTorch：

1. # -*- coding: utf-8 -*-
   
2. """
   
3. Created on Sun Feb  5 14:56:26 2023
   
4. @author: halcom
   
5. """
   
6. import sys
   
7. import os
   
8. # os.environ['CUDA_VISIBLE_DEVICES'] = '0'
   
9. import numpy as np
   
10. import torch
    
11. import torch.utils.data as data
    
12. from torch.utils.data import DataLoader
    
13. from torch import autograd, optim
    
14. import torch.nn as nn
    
15. import torch.nn.functional as F
    
16. from torch.nn.modules.loss import _Loss
    
17. from torchvision.transforms import transforms
    
18. from functools import partial
    
19. from xml.dom.minidom import parse
    
20. import cv2
    
21. import time
    
22. import shutil
    
23. import copy
    
24. from efficientnet_pytorch import EfficientNet
    
25. #
    
26. # 将字符串转化为数字标签，序号从1开始，1-N
    
27. class strLabelConverter(object):
    
28.     """Convert between str and label.
    
29.     NOTE:
    
30.         Insert `blank` to the alphabet for CTC.
    
31.     Args:
    
32.         alphabet (str): set of the possible characters.
    
33.         ignore_case (bool, default=True): whether or not to ignore all of the case.
    
34.     """
    
35. 
    
36.     def __init__(self, alphabet, ignore_case=False):
    
37.         self._ignore_case = ignore_case
    
38.         if self._ignore_case:
    
39.             alphabet = alphabet.lower()
    
40.         self.alphabet = alphabet + '-'  # for `-1` index
    
41. 
    
42.         self.dict = {}
    
43.         for i, char in enumerate(alphabet):
    
44.             # NOTE: 0 is reserved for 'blank' required by wrap_ctc
    
45.             self.dict[char] = i + 1
    
46. 
    
47.     def encode(self, text):
    
48.         """Support batch or single str.
    
49.         Args:
    
50.             text (str or list of str): texts to convert.
    
51.         Returns:
    
52.             torch.IntTensor [length_0 + length_1 + ... length_{n - 1}]: encoded texts.
    
53.             torch.IntTensor [n]: length of each text.
    
54.         """
    
55. 
    
56.         length = []
    
57.         result = []
    
58.         decode_flag = True if type(text[0])==bytes else False
    
59. 
    
60.         for item in text:
    
61. 
    
62.             if decode_flag:
    
63.                 item = item.decode('utf-8','strict')
    
64.             length.append(len(item))
    
65.             for char in item:
    
66.                 index = self.dict[char]
    
67.                 result.append(index)
    
68.         text = result
    
69.         return (torch.IntTensor(text), torch.IntTensor(length))
    
70. 
    
71.     def decode(self, t, length, raw=False):
    
72.         """Decode encoded texts back into strs.
    
73.         Args:
    
74.             torch.IntTensor [length_0 + length_1 + ... length_{n - 1}]: encoded texts.
    
75.             torch.IntTensor [n]: length of each text.
    
76.         Raises:
    
77.             AssertionError: when the texts and its length does not match.
    
78.         Returns:
    
79.             text (str or list of str): texts to convert.
    
80.         """
    
81.         if length.numel() == 1:
    
82.             length = length[0]
    
83.             assert t.numel() == length, "text with length: {} does not match declared length: {}".format(t.numel(), length)
    
84.             if raw:
    
85.                 return ''.join([self.alphabet[i - 1] for i in t])
    
86.             else:
    
87.                 char_list = []
    
88.                 for i in range(length):
    
89.                     if t[i] != 0 and (not (i > 0 and t[i - 1] == t[i])):
    
90.                         char_list.append(self.alphabet[t[i] - 1])
    
91.                 return ''.join(char_list)
    
92.         else:
    
93.             # batch mode
    
94.             assert t.numel() == length.sum(), "texts with length: {} does not match declared length: {}".format(t.numel(), length.sum())
    
95.             texts = []
    
96.             index = 0
    
97.             for i in range(length.numel()):
    
98.                 l = length[i]
    
99.                 texts.append(
    
100.                     self.decode(
     
101.                         t[index:index + l], torch.IntTensor([l]), raw=raw))
     
102.                 index += l
     
103.             return texts
     
104. #
     
105. # 测试函数strLabelConverter
     
106. def AlphaConverter(AlphaBet):
     
107.     AlphaBet = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
     
108.     converter = strLabelConverter(AlphaBet)
     
109.     labels = '0134ABC' # 注意 CrossEntropyLoss， label需要从0到N-1
     
110.     text_train, length_train = converter.encode(labels)
     
111.     # calc correct ratio
     
112.     preds = torch.tensor([1, 3, 4, 11])
     
113.     preds_size = torch.IntTensor([preds.size(0)])
     
114.     sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
     
115.     # sim_preds对应的就是AlphaBet
     
116.     return converter
     
117. #
     
118. def make_dataset_train_txt(root, txt_file, train_image_format):
     
119.     imgs=[]
     
120.     with open(txt_file) as f:
     
121.         indexs=f.readlines()
     
122.     for i in range( len(indexs) ):
     
123. #        index = str( indexs[i] ).split('.')[0]
     
124.         index = str( indexs[i] ).split('\n')[0]
     
125.         # img=os.path.join(root, "train", index+".png")
     
126.         img=os.path.join(root, "Image", index + train_image_format)
     
127.         mask=os.path.join(root, "Mask", index + ".png")
     
128.         imgs.append((img,mask))
     
129.     return imgs
     
130. def make_dataset_val_txt(root, txt_file, train_image_format):
     
131.     imgs=[]
     
132.     with open(txt_file) as f:
     
133.         indexs=f.readlines()
     
134.     for i in range( len(indexs) ):
     
135. #        index = str( indexs[i] ).split('.')[0]
     
136.         index = str( indexs[i] ).split('\n')[0]
     
137.         # img=os.path.join(root, "val", index+".png")
     
138.         img=os.path.join(root, "Image", index + train_image_format)
     
139.         mask=os.path.join(root, "Mask", index + ".png")
     
140.         imgs.append((img,mask))
     
141.     return imgs
     
142. class train_LiverDataset(data.Dataset):
     
143.     def __init__(self, root, txt_file, classNum = 1, transform = None, target_transform = None, train_image_format = ".jpg", img_channel = 3):
     
144.         imgs = make_dataset_train_txt(root, txt_file, train_image_format)
     
145.         self.imgs = imgs
     
146.         self.transform = transform
     
147.         self.target_transform = target_transform
     
148.         self.classNum = classNum
     
149.         self.train_image_format = train_image_format
     
150.         self.img_channel = img_channel
     
151.     def __getitem__(self, index):
     
152.         x_path, y_path = self.imgs[index]
     
153.         if(self.img_channel == 3):
     
154.             image = cv2.imread(x_path, 1)
     
155.             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
156.         elif(self.img_channel == 1):
     
157.             image = cv2.imread(x_path, 0)
     
158.         else:
     
159.             return
     
160.         mask = cv2.imread(y_path, 0)
     
161.         # extract certain classes from mask
     
162. 
     
163.         mask = np.stack(masks, axis=-1).astype('float')
     
164.         
     
165.         if self.transform is not None:
     
166.             image = self.transform(image)
     
167.         if self.target_transform is not None:
     
168.             mask = self.target_transform(mask)
     
169.         return image, mask
     
170.     def __len__(self):
     
171.         return len(self.imgs)
     
172. class val_LiverDataset(data.Dataset):
     
173.     def __init__(self, root, txt_file, classNum = 1, transform = None, target_transform = None, train_image_format = ".jpg", img_channel = 3):
     
174.         imgs = make_dataset_val_txt(root, txt_file, train_image_format)
     
175.         self.imgs = imgs
     
176.         self.transform = transform
     
177.         self.target_transform = target_transform
     
178.         self.classNum = classNum
     
179.         self.train_image_format = train_image_format
     
180.         self.img_channel = img_channel
     
181.     def __getitem__(self, index):
     
182.         x_path, y_path = self.imgs[index]
     
183.         if(self.img_channel == 3):
     
184.             image = cv2.imread(x_path, 1)
     
185.             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
186.         elif(self.img_channel == 1):
     
187.             image = cv2.imread(x_path, 0)
     
188.         else:
     
189.             return
     
190.         mask = cv2.imread(y_path, 0)
     
191.         # extract certain classes from mask
     
192. 
     
193.         mask = np.stack(masks, axis=-1).astype('float')
     
194.         
     
195.         if self.transform is not None:
     
196.             image = self.transform(image)
     
197.         if self.target_transform is not None:
     
198.             mask = self.target_transform(mask)
     
199.         return image, mask
     
200.     def __len__(self):
     
201.         return len(self.imgs)
     
202. # training model
     
203. def trainAI(train_project_path, xmlfile, img_root_path, mask_root_path, train_txt_file_path, train_loss_txt_file, val_loss_txt_file, trainbatch_loss_txt_file, train_image_format, initial_model, GPU_index, img_channel, img_size, batch_size, num_workers, num_epochs, Seg_p, lr_init, moment_init, classNum, ENCODER, ENCODER_WEIGHTS, AlphaBet):
     
204.     # 是否使用cuda
     
205.     device = torch.device("cuda:" + str(GPU_index) if torch.cuda.is_available() else "cpu")
     
206.     classNum = len(AlphaBet)
     
207.     model = EfficientNet.from_pretrained('efficientnet-b0', in_channels=img_channel, num_classes= len(AlphaBet)).to(device)
     
208.     criterion = torch.nn.CrossEntropyLoss().to(device)
     
209.     optimizer = torch.optim.Adam([dict(params=model.parameters(), lr = lr_init)])
     
210.     if img_channel == 1:
     
211.         x_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) ])
     
212.     elif img_channel == 3:
     
213.         x_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ])
     
214.         # x_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])
     
215.     else:
     
216.         return
     
217.     y_transforms = transforms.ToTensor()
     
218.     train_liver_dataset = train_LiverDataset(train_txt_file_path, train_txt_file_path + "/Mask_train.txt", classNum = classNum, transform = x_transforms, target_transform = y_transforms, train_image_format = train_image_format, img_channel = img_channel)
     
219.     train_dataloaders = DataLoader(train_liver_dataset, batch_size = batch_size, shuffle = True, num_workers = num_workers)
     
220.     val_liver_dataset = val_LiverDataset(train_txt_file_path, train_txt_file_path + "/Mask_validation.txt", classNum = classNum, transform = x_transforms, target_transform = y_transforms, train_image_format = train_image_format, img_channel = img_channel)
     
221.     val_dataloaders = DataLoader(val_liver_dataset, batch_size = batch_size, shuffle = True, num_workers = num_workers)
     
222.     if initial_model != "None":
     
223.         model = torch.jit.load(initial_model)
     
224.     model.cuda(GPU_index)
     
225.     model.train()
     
226.     train_sum_loss = np.inf
     
227.     val_sum_loss = np.inf
     
228.     for epoch in range(num_epochs):
     
229.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
     
230.         dt_size = len(train_dataloaders.dataset)
     
231.         step = 0
     
232.         epoch_loss = 0
     
233.         model.train()
     
234.         for x, y in train_dataloaders:
     
235.             step += 1
     
236.             inputs = x.to(device)
     
237.             labels = y.to(device)
     
238.             optimizer.zero_grad()
     
239.             outputs = model(inputs)
     
240.             loss_dict = criterion(outputs, labels)
     
241.             loss.backward()
     
242.             optimizer.step()
     
243.             epoch_loss += loss.item()
     
244. 
     
245.             print("%d/%d,train_loss:%0.3f,train_metric:%0.3f" % (step, (dt_size - 1) // train_dataloaders.batch_size + 1, loss.item(), epoch_metric))
     
246.             del x, y, inputs, labels, outputs
     
247.             torch.cuda.empty_cache()
     
248.         # print("epoch %d loss:%0.3f" % (epoch, epoch_loss))
     
249.         if(epoch_loss < train_sum_loss):
     
250.             train_sum_loss = epoch_loss
     
251.             model.eval()
     
252.             example = torch.rand(1, img_channel, img_size, img_size).type(torch.FloatTensor).cuda(GPU_index)
     
253.             pt_model = torch.jit.trace(model, example)
     
254.             pt_model.save( os.path.join(train_project_path, "Model", str(epoch) +'-Loss.pt') )

复制代码

参考：
【1】https://github.com/lukemelas/EfficientNet-PyTorch
【2】百度网盘连接：链接：https://pan.baidu.com/s/1dVcTded9PD-8Oy270bAGhA 提取码：u60r
【3】Efficientnet_pytorch训练&#12163;&#12080;的数据集，并对数据进&#12175;分类
【4】【图像分类】——来来来，干了这碗EfficientNet实战（Pytorch）

Halcom

1. 
   
2. efficientnet              1.1.1                    pypi_0    pypi
   
3. efficientnet-pytorch      0.6.3                    pypi_0    pypi
   
4. 
   
5. torch                     1.4.0                    pypi_0    pypi
   
6. torchvision               0.5.0                    pypi_0    pypi
   
7. 
   
8. python                    3.7.6                h60c2a47_2
   
9. python-dateutil           2.8.1                      py_0
   
10. python-jsonrpc-server     0.3.4                      py_0
    
11. python-language-server    0.31.7                   py37_0
    
12. python-libarchive-c       2.8                     py37_13
    
13. 
    
14. opencv-python             4.4.0.46                 pypi_0    pypi
    
15. openpyxl                  3.0.3                      py_0
    
16. openssl                   1.1.1d               he774522_4

复制代码