Upload 6 files
Browse files- .gitattributes +1 -0
- characters.zip +3 -0
- dataset.py +80 -0
- prepare_data.py +100 -0
- simhei.ttf +3 -0
- train.py +353 -0
- 形近字.txt +1 -0
.gitattributes
CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
36 |
+
simhei.ttf filter=lfs diff=lfs merge=lfs -text
|
characters.zip
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:d2599ed6416f8272ed2bfa803410db6f4b0bfa8ff5b76c3b949ca50842688604
|
3 |
+
size 17976013
|
dataset.py
ADDED
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from torchvision.datasets import MNIST
|
2 |
+
import os
|
3 |
+
import numpy as np
|
4 |
+
import random
|
5 |
+
|
6 |
+
train_dataset = MNIST(os.path.join('./', "MNIST"), train=True, download=True)
|
7 |
+
test_dataset = MNIST(os.path.join('./', "MNIST"), train=False, download=True)
|
8 |
+
|
9 |
+
|
10 |
+
class MNIST_DS(object):
|
11 |
+
|
12 |
+
def __init__(self, train_dataset, test_dataset):
|
13 |
+
self.__train_labels_idx_map = {}
|
14 |
+
self.__test_labels_idx_map = {}
|
15 |
+
|
16 |
+
self.__train_data = train_dataset.data
|
17 |
+
self.__test_data = test_dataset.data
|
18 |
+
self.__train_labels = train_dataset.targets
|
19 |
+
self.__test_labels = test_dataset.targets
|
20 |
+
|
21 |
+
self.__train_labels_np = self.__train_labels.numpy()
|
22 |
+
self.__train_unique_labels = np.unique(self.__train_labels_np)
|
23 |
+
|
24 |
+
self.__test_labels_np = self.__test_labels.numpy()
|
25 |
+
self.__test_unique_labels = np.unique(self.__test_labels_np)
|
26 |
+
|
27 |
+
def load(self):
|
28 |
+
self.__train_labels_idx_map = {}
|
29 |
+
for label in self.__train_unique_labels:
|
30 |
+
self.__train_labels_idx_map[label] = np.where(self.__train_labels_np == label)[0]
|
31 |
+
|
32 |
+
self.__test_labels_idx_map = {}
|
33 |
+
for label in self.__test_unique_labels:
|
34 |
+
self.__test_labels_idx_map[label] = np.where(self.__test_labels_np == label)[0]
|
35 |
+
|
36 |
+
def getTriplet(self, split="train"):
|
37 |
+
pos_label = 0
|
38 |
+
neg_label = 0
|
39 |
+
label_idx_map = None
|
40 |
+
data = None
|
41 |
+
|
42 |
+
if split == 'train':
|
43 |
+
pos_label = self.__train_unique_labels[random.randint(0, len(self.__train_unique_labels) - 1)]
|
44 |
+
neg_label = pos_label
|
45 |
+
while neg_label is pos_label:
|
46 |
+
neg_label = self.__train_unique_labels[random.randint(0, len(self.__train_unique_labels) - 1)]
|
47 |
+
label_idx_map = self.__train_labels_idx_map
|
48 |
+
data = self.__train_data
|
49 |
+
else:
|
50 |
+
pos_label = self.__test_unique_labels[random.randint(0, len(self.__test_unique_labels) - 1)]
|
51 |
+
neg_label = pos_label
|
52 |
+
while neg_label is pos_label:
|
53 |
+
neg_label = self.__test_unique_labels[random.randint(0, len(self.__test_unique_labels) - 1)]
|
54 |
+
label_idx_map = self.__test_labels_idx_map
|
55 |
+
data = self.__test_data
|
56 |
+
|
57 |
+
pos_label_idx_map = label_idx_map[pos_label]
|
58 |
+
pos_img_anchor_idx = pos_label_idx_map[random.randint(0, len(pos_label_idx_map) - 1)]
|
59 |
+
pos_img_idx = pos_img_anchor_idx
|
60 |
+
while pos_img_idx is pos_img_anchor_idx:
|
61 |
+
pos_img_idx = pos_label_idx_map[random.randint(0, len(pos_label_idx_map) - 1)]
|
62 |
+
|
63 |
+
neg_label_idx_map = label_idx_map[neg_label]
|
64 |
+
neg_img_idx = neg_label_idx_map[random.randint(0, len(neg_label_idx_map) - 1)]
|
65 |
+
|
66 |
+
pos_anchor_img = data[pos_img_anchor_idx].numpy()
|
67 |
+
pos_img = data[pos_img_idx].numpy()
|
68 |
+
neg_img = data[neg_img_idx].numpy()
|
69 |
+
|
70 |
+
return pos_anchor_img, pos_img, neg_img
|
71 |
+
|
72 |
+
|
73 |
+
dset_obj = MNIST_DS(train_dataset, test_dataset)
|
74 |
+
dset_obj.load()
|
75 |
+
train_triplets = []
|
76 |
+
pos_anchor_img, pos_img, neg_img = dset_obj.getTriplet()
|
77 |
+
train_triplets.append([pos_anchor_img, pos_img, neg_img])
|
78 |
+
|
79 |
+
print(train_triplets[0][0])
|
80 |
+
|
prepare_data.py
ADDED
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding: utf-8 -*-
|
2 |
+
import random
|
3 |
+
import argparse
|
4 |
+
|
5 |
+
import pygame
|
6 |
+
from PIL import Image
|
7 |
+
import numpy as np
|
8 |
+
import os
|
9 |
+
from pypinyin import lazy_pinyin
|
10 |
+
|
11 |
+
|
12 |
+
def set_seed():
|
13 |
+
random.seed(args.seed)
|
14 |
+
|
15 |
+
|
16 |
+
def load_chars(path):
|
17 |
+
# 获取所有汉字列表
|
18 |
+
chars_list = []
|
19 |
+
with open(path, 'r+', encoding='utf-8') as f:
|
20 |
+
lines = f.readlines()
|
21 |
+
for line in lines:
|
22 |
+
chars = [char for char in line.strip()]
|
23 |
+
chars_list.append(chars)
|
24 |
+
|
25 |
+
return chars_list
|
26 |
+
|
27 |
+
|
28 |
+
def char2image(chars_list):
|
29 |
+
# 通过pygame将汉字转化为黑白图片
|
30 |
+
pygame.init()
|
31 |
+
output_dir_path = os.path.join('./', args.output_path)
|
32 |
+
if not os.path.exists(output_dir_path):
|
33 |
+
os.makedirs(output_dir_path)
|
34 |
+
|
35 |
+
train_ids_list, test_ids_list = data_split(len(chars_list), args.train_percentage)
|
36 |
+
|
37 |
+
for i, chars in enumerate(chars_list):
|
38 |
+
if i in train_ids_list:
|
39 |
+
chars_class_path = os.path.join(output_dir_path, 'train', f'character_{i + 1}')
|
40 |
+
elif i in test_ids_list:
|
41 |
+
chars_class_path = os.path.join(output_dir_path, 'test', f'character_{i + 1}')
|
42 |
+
else:
|
43 |
+
raise ValueError(f"The length of dataset is out of range, which idx is {i}")
|
44 |
+
|
45 |
+
if not os.path.exists(chars_class_path):
|
46 |
+
os.makedirs(chars_class_path)
|
47 |
+
|
48 |
+
for j, char in enumerate(chars):
|
49 |
+
char_path = os.path.join(chars_class_path, f'{i + 1}_{j + 1}.png')
|
50 |
+
if os.path.exists(char_path):
|
51 |
+
continue
|
52 |
+
# 文件夹里还有别的类型的字体
|
53 |
+
font = pygame.font.Font(r"C://Windows/Fonts/simhei.ttf", 100)
|
54 |
+
# 第三个参数为字体颜色,第四个参数为背景颜色。
|
55 |
+
rtext = font.render(char, True, (0, 0, 0), (255, 255, 255))
|
56 |
+
pygame.image.save(rtext, char_path)
|
57 |
+
|
58 |
+
|
59 |
+
def data_split(num_examples, train_percentage):
|
60 |
+
set_seed()
|
61 |
+
num_train_examples = int(num_examples * train_percentage)
|
62 |
+
num_test_examples = num_examples - num_train_examples
|
63 |
+
|
64 |
+
train_dir_path = os.path.join(args.output_path, 'train')
|
65 |
+
test_dir_path = os.path.join(args.output_path, 'test')
|
66 |
+
if not os.path.exists(train_dir_path):
|
67 |
+
os.makedirs(train_dir_path)
|
68 |
+
if not os.path.exists(test_dir_path):
|
69 |
+
os.makedirs(test_dir_path)
|
70 |
+
|
71 |
+
train_ids_list = random.sample(range(num_examples), num_train_examples)
|
72 |
+
test_ids_list = [idx for idx in range(num_examples) if idx not in train_ids_list]
|
73 |
+
|
74 |
+
assert len(test_ids_list) == num_test_examples
|
75 |
+
return train_ids_list, test_ids_list
|
76 |
+
|
77 |
+
|
78 |
+
def main():
|
79 |
+
path = './形近字.txt'
|
80 |
+
chars_list = load_chars(path)
|
81 |
+
num_examples = len(chars_list)
|
82 |
+
char2image(chars_list)
|
83 |
+
|
84 |
+
|
85 |
+
if __name__ == '__main__':
|
86 |
+
# TODO:其他类型的字体
|
87 |
+
parser = argparse.ArgumentParser(description='Make the character-pairs dataset')
|
88 |
+
parser.add_argument('--output_path', default='characters', type=str,
|
89 |
+
help='Path to store dataset')
|
90 |
+
parser.add_argument('--seed', default=718, type=int,
|
91 |
+
help='Fixed seed for Random package')
|
92 |
+
parser.add_argument('--dataset_path', default='形近字.txt', type=str,
|
93 |
+
help='Path of raw dataset')
|
94 |
+
parser.add_argument('--train_percentage', default=0.7, type=float,
|
95 |
+
help='Percentage of training set')
|
96 |
+
|
97 |
+
global args
|
98 |
+
args = parser.parse_args()
|
99 |
+
|
100 |
+
main()
|
simhei.ttf
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:aa4560dd8fe5645745fed3ffa301c3ca4d6c03cbd738145b613303961ba733b8
|
3 |
+
size 9753388
|
train.py
ADDED
@@ -0,0 +1,353 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
import torch.optim as optim
|
4 |
+
import torch.nn.functional as F
|
5 |
+
import os
|
6 |
+
import random
|
7 |
+
import argparse
|
8 |
+
from torchvision import transforms
|
9 |
+
from torch.autograd import Variable
|
10 |
+
import cv2
|
11 |
+
import numpy as np
|
12 |
+
|
13 |
+
|
14 |
+
class BaseLoader(torch.utils.data.Dataset):
|
15 |
+
def __init__(self, triplets, transform=None):
|
16 |
+
self.triplets = triplets
|
17 |
+
self.transform = transform
|
18 |
+
|
19 |
+
def __getitem__(self, index):
|
20 |
+
img1_pth, img2_pth, img3_pth = self.triplets[index]
|
21 |
+
img1 = cv2.imread(img1_pth)
|
22 |
+
img2 = cv2.imread(img2_pth)
|
23 |
+
img3 = cv2.imread(img3_pth)
|
24 |
+
|
25 |
+
try:
|
26 |
+
img1 = cv2.resize(img1, (args.picture_resize, args.picture_resize))
|
27 |
+
except Exception as e:
|
28 |
+
img1 = np.zeros((args.picture_resize, args.picture_resize, 3), dtype=np.uint8)
|
29 |
+
|
30 |
+
try:
|
31 |
+
img2 = cv2.resize(img2, (args.picture_resize, args.picture_resize))
|
32 |
+
except Exception as e:
|
33 |
+
img2 = np.zeros((args.picture_resize, args.picture_resize, 3), dtype=np.uint8)
|
34 |
+
|
35 |
+
try:
|
36 |
+
img3 = cv2.resize(img3, (args.picture_resize, args.picture_resize))
|
37 |
+
except Exception as e:
|
38 |
+
img3 = np.zeros((args.picture_resize, args.picture_resize, 3), dtype=np.uint8)
|
39 |
+
|
40 |
+
if self.transform is not None:
|
41 |
+
img1 = self.transform(img1)
|
42 |
+
img2 = self.transform(img2)
|
43 |
+
img3 = self.transform(img3)
|
44 |
+
|
45 |
+
return img1, img2, img3
|
46 |
+
|
47 |
+
def __len__(self):
|
48 |
+
return len(self.triplets)
|
49 |
+
|
50 |
+
|
51 |
+
class BaseCnn(nn.Module):
|
52 |
+
def __init__(self):
|
53 |
+
super(BaseCnn, self).__init__()
|
54 |
+
self.conv1 = nn.Sequential(
|
55 |
+
nn.Conv2d(3, 64, 3),
|
56 |
+
nn.MaxPool2d(2)
|
57 |
+
)
|
58 |
+
self.conv2 = nn.Sequential(
|
59 |
+
nn.Conv2d(64, 128, 3),
|
60 |
+
nn.MaxPool2d(2)
|
61 |
+
)
|
62 |
+
self.conv3 = nn.Sequential(
|
63 |
+
nn.Conv2d(128, 128, 3),
|
64 |
+
nn.MaxPool2d(2)
|
65 |
+
)
|
66 |
+
|
67 |
+
def forward(self, x):
|
68 |
+
x = self.conv1(x)
|
69 |
+
x = self.conv2(x)
|
70 |
+
x = self.conv3(x)
|
71 |
+
x = x.view(x.size(0), -1)
|
72 |
+
x = F.normalize(x, p=2, dim=1)
|
73 |
+
return x
|
74 |
+
|
75 |
+
|
76 |
+
class SiameseNet(nn.Module):
|
77 |
+
def __init__(self):
|
78 |
+
super(SiameseNet, self).__init__()
|
79 |
+
self.base = BaseCnn()
|
80 |
+
|
81 |
+
def forward(self, x1, x2, x3):
|
82 |
+
x1 = self.base(x1)
|
83 |
+
x2 = self.base(x2)
|
84 |
+
x3 = self.base(x3)
|
85 |
+
|
86 |
+
return x1, x2, x3
|
87 |
+
|
88 |
+
|
89 |
+
class BaseDset(object):
|
90 |
+
def __init__(self):
|
91 |
+
self.__base_path = ""
|
92 |
+
|
93 |
+
self.__train_set = {}
|
94 |
+
self.__test_set = {}
|
95 |
+
self.__train_keys = []
|
96 |
+
self.__test_keys = []
|
97 |
+
|
98 |
+
def load(self, base_path):
|
99 |
+
"""加载数据集,将类别和路径存储"""
|
100 |
+
self.__base_path = base_path
|
101 |
+
train_dir = os.path.join(self.__base_path, 'train')
|
102 |
+
test_dir = os.path.join(self.__base_path, 'test')
|
103 |
+
|
104 |
+
self.__train_set = {}
|
105 |
+
self.__test_set = {}
|
106 |
+
self.__train_keys = []
|
107 |
+
self.__test_keys = []
|
108 |
+
|
109 |
+
for class_id in os.listdir(train_dir):
|
110 |
+
# 对于train_dir里的每个文件夹名字 classi
|
111 |
+
class_dir = os.path.join(train_dir, class_id)
|
112 |
+
# 为其在训练集合中创建一个文件夹
|
113 |
+
# 在类别集合中,即train_keys中添加类别classi
|
114 |
+
self.__train_set[class_id] = []
|
115 |
+
self.__train_keys.append(class_id)
|
116 |
+
# 对于每个类别内的数据,将其路径添加到集合中
|
117 |
+
for img_name in os.listdir(class_dir):
|
118 |
+
img_path = os.path.join(class_dir, img_name)
|
119 |
+
self.__train_set[class_id].append(img_path)
|
120 |
+
# 同理对于测试集合也一样
|
121 |
+
for class_id in os.listdir(test_dir):
|
122 |
+
class_dir = os.path.join(test_dir, class_id)
|
123 |
+
self.__test_set[class_id] = []
|
124 |
+
self.__test_keys.append(class_id)
|
125 |
+
for img_name in os.listdir(class_dir):
|
126 |
+
img_path = os.path.join(class_dir, img_name)
|
127 |
+
self.__test_set[class_id].append(img_path)
|
128 |
+
|
129 |
+
return len(self.__train_keys), len(self.__test_keys)
|
130 |
+
|
131 |
+
# 获取三元组 !!!
|
132 |
+
def getTriplet(self, split='train'):
|
133 |
+
# 默认选取训练集
|
134 |
+
if split == 'train':
|
135 |
+
dataset = self.__train_set
|
136 |
+
keys = self.__train_keys
|
137 |
+
else:
|
138 |
+
dataset = self.__test_set
|
139 |
+
keys = self.__test_keys
|
140 |
+
|
141 |
+
# 随机指定两个正负类别,确保二者不一致
|
142 |
+
pos_idx = random.randint(0, len(keys) - 1)
|
143 |
+
while True:
|
144 |
+
neg_idx = random.randint(0, len(keys) - 1)
|
145 |
+
if pos_idx != neg_idx:
|
146 |
+
break
|
147 |
+
# 选定一个原始样本
|
148 |
+
pos_anchor_img_idx = random.randint(0, len(dataset[keys[pos_idx]]) - 1)
|
149 |
+
# 随机选择一个正样本,保证二者不一致
|
150 |
+
while True:
|
151 |
+
pos_img_idx = random.randint(0, len(dataset[keys[pos_idx]]) - 1)
|
152 |
+
if pos_anchor_img_idx != pos_img_idx:
|
153 |
+
break
|
154 |
+
# 随机选择一个负样本
|
155 |
+
neg_img_idx = random.randint(0, len(dataset[keys[neg_idx]]) - 1)
|
156 |
+
|
157 |
+
# 生成三元组
|
158 |
+
pos_anchor_img = dataset[keys[pos_idx]][pos_anchor_img_idx]
|
159 |
+
pos_img = dataset[keys[pos_idx]][pos_img_idx]
|
160 |
+
neg_img = dataset[keys[neg_idx]][neg_img_idx]
|
161 |
+
|
162 |
+
return pos_anchor_img, pos_img, neg_img
|
163 |
+
|
164 |
+
|
165 |
+
def train(data, model, criterion, optimizer, epoch):
|
166 |
+
print("******** Training ********")
|
167 |
+
total_loss = 0
|
168 |
+
model.train()
|
169 |
+
for batch_idx, img_triplet in enumerate(data):
|
170 |
+
# 提取数据
|
171 |
+
anchor_img, pos_img, neg_img = img_triplet
|
172 |
+
anchor_img, pos_img, neg_img = anchor_img.to(device), pos_img.to(device), neg_img.to(device)
|
173 |
+
anchor_img, pos_img, neg_img = Variable(anchor_img), Variable(pos_img), Variable(neg_img)
|
174 |
+
# 分别获得三个编码
|
175 |
+
E1, E2, E3 = model(anchor_img, pos_img, neg_img)
|
176 |
+
# 计算二者之间的欧式距离
|
177 |
+
dist_E1_E2 = F.pairwise_distance(E1, E2, 2)
|
178 |
+
dist_E1_E3 = F.pairwise_distance(E1, E3, 2)
|
179 |
+
|
180 |
+
target = torch.FloatTensor(dist_E1_E2.size()).fill_(-1)
|
181 |
+
target = target.to(device)
|
182 |
+
target = Variable(target)
|
183 |
+
# 大小如何?
|
184 |
+
loss = criterion(dist_E1_E2, dist_E1_E3, target)
|
185 |
+
total_loss += loss
|
186 |
+
|
187 |
+
optimizer.zero_grad()
|
188 |
+
loss.backward()
|
189 |
+
optimizer.step()
|
190 |
+
# 打印一波损失
|
191 |
+
log_step = args.train_log_step
|
192 |
+
if (batch_idx % log_step == 0) and (batch_idx != 0):
|
193 |
+
print('Train Epoch: {} [{}/{}] \t Loss: {:.4f}'.format(epoch, batch_idx, len(data), total_loss / log_step))
|
194 |
+
total_loss = 0
|
195 |
+
print("****************")
|
196 |
+
|
197 |
+
|
198 |
+
def test(data, model, criterion):
|
199 |
+
print("******** Testing ********")
|
200 |
+
with torch.no_grad():
|
201 |
+
model.eval()
|
202 |
+
accuracies = [0, 0, 0]
|
203 |
+
acc_threshes = [0, 0.2, 0.5]
|
204 |
+
total_loss = 0
|
205 |
+
for batch_idx, img_triplet in enumerate(data):
|
206 |
+
anchor_img, pos_img, neg_img = img_triplet
|
207 |
+
anchor_img, pos_img, neg_img = anchor_img.to(device), pos_img.to(device), neg_img.to(device)
|
208 |
+
anchor_img, pos_img, neg_img = Variable(anchor_img), Variable(pos_img), Variable(neg_img)
|
209 |
+
E1, E2, E3 = model(anchor_img, pos_img, neg_img)
|
210 |
+
dist_E1_E2 = F.pairwise_distance(E1, E2, 2)
|
211 |
+
dist_E1_E3 = F.pairwise_distance(E1, E3, 2)
|
212 |
+
|
213 |
+
target = torch.FloatTensor(dist_E1_E2.size()).fill_(-1)
|
214 |
+
target = target.to(device)
|
215 |
+
target = Variable(target)
|
216 |
+
|
217 |
+
loss = criterion(dist_E1_E2, dist_E1_E3, target)
|
218 |
+
total_loss += loss
|
219 |
+
|
220 |
+
for i in range(len(accuracies)):
|
221 |
+
prediction = (dist_E1_E3 - dist_E1_E2 - args.margin * acc_threshes[i]).cpu().data
|
222 |
+
prediction = prediction.view(prediction.numel())
|
223 |
+
prediction = (prediction > 0).float()
|
224 |
+
batch_acc = prediction.sum() * 1.0 / prediction.numel()
|
225 |
+
accuracies[i] += batch_acc
|
226 |
+
print('Test Loss: {}'.format(total_loss / len(data)))
|
227 |
+
for i in range(len(accuracies)):
|
228 |
+
# 0%等价于准确率其余是更严格的指标
|
229 |
+
print(
|
230 |
+
'Test Accuracy with diff = {}% of margin: {:.4f}'.format(acc_threshes[i] * 100,
|
231 |
+
accuracies[i] / len(data)))
|
232 |
+
print("****************")
|
233 |
+
|
234 |
+
return accuracies[-1]
|
235 |
+
|
236 |
+
|
237 |
+
def main():
|
238 |
+
# random_seed
|
239 |
+
torch.manual_seed(718)
|
240 |
+
torch.cuda.manual_seed(718)
|
241 |
+
|
242 |
+
data_path = r'./characters'
|
243 |
+
# data_path = r'./characters'
|
244 |
+
dset_obj = BaseDset()
|
245 |
+
dset_obj.load(data_path)
|
246 |
+
|
247 |
+
train_triplets = []
|
248 |
+
test_triplets = []
|
249 |
+
|
250 |
+
for i in range(args.num_train_samples):
|
251 |
+
pos_anchor_img, pos_img, neg_img = dset_obj.getTriplet()
|
252 |
+
train_triplets.append([pos_anchor_img, pos_img, neg_img])
|
253 |
+
for i in range(args.num_test_samples):
|
254 |
+
pos_anchor_img, pos_img, neg_img = dset_obj.getTriplet(split='test')
|
255 |
+
test_triplets.append([pos_anchor_img, pos_img, neg_img])
|
256 |
+
loader = BaseLoader
|
257 |
+
model = SiameseNet()
|
258 |
+
model.to(device)
|
259 |
+
|
260 |
+
criterion = torch.nn.MarginRankingLoss(margin=args.margin)
|
261 |
+
optimizer = optim.Adam(model.parameters(), lr=args.lr)
|
262 |
+
|
263 |
+
best_acc_of_50_margin = 0
|
264 |
+
best_epoch = 0
|
265 |
+
|
266 |
+
for epoch in range(1, args.epochs + 1):
|
267 |
+
# 初始化数据加载器
|
268 |
+
# 加载三元组
|
269 |
+
train_data_loader = torch.utils.data.DataLoader(
|
270 |
+
loader(train_triplets,
|
271 |
+
transform=transforms.Compose([
|
272 |
+
transforms.ToTensor(),
|
273 |
+
transforms.Normalize(0, 1)
|
274 |
+
])),
|
275 |
+
batch_size=args.batch_size, shuffle=True)
|
276 |
+
test_data_loader = torch.utils.data.DataLoader(
|
277 |
+
loader(test_triplets,
|
278 |
+
transform=transforms.Compose([
|
279 |
+
transforms.ToTensor(),
|
280 |
+
transforms.Normalize(0, 1)
|
281 |
+
])),
|
282 |
+
batch_size=args.batch_size, shuffle=True)
|
283 |
+
train(train_data_loader, model, criterion, optimizer, epoch)
|
284 |
+
acc_of_50_margin = test(test_data_loader, model, criterion)
|
285 |
+
|
286 |
+
model_to_save = {
|
287 |
+
"epoch": epoch + 1,
|
288 |
+
'state_dict': model.state_dict(),
|
289 |
+
}
|
290 |
+
|
291 |
+
if acc_of_50_margin > best_acc_of_50_margin:
|
292 |
+
best_acc_of_50_margin = acc_of_50_margin
|
293 |
+
best_epoch = epoch
|
294 |
+
|
295 |
+
if not args.disable_save_best_ckp:
|
296 |
+
result_path = os.path.join(args.result_dir)
|
297 |
+
file_name = os.path.join(args.result_dir, "best_checkpoint" + ".pt")
|
298 |
+
if not os.path.exists(result_path):
|
299 |
+
os.makedirs(result_path)
|
300 |
+
save_checkpoint(model_to_save, file_name)
|
301 |
+
|
302 |
+
if (epoch % args.ckp_freq == 0) and not args.disable_save_ckp:
|
303 |
+
result_path = os.path.join(args.result_dir)
|
304 |
+
file_name = os.path.join(args.result_dir, "checkpoint_" + str(epoch) + ".pt")
|
305 |
+
if not os.path.exists(result_path):
|
306 |
+
os.makedirs(result_path)
|
307 |
+
save_checkpoint(model_to_save, file_name)
|
308 |
+
print("Training is done.")
|
309 |
+
print(f"The best epoch of acc50, which is {best_acc_of_50_margin * 100}%, is {best_epoch}.")
|
310 |
+
|
311 |
+
|
312 |
+
def save_checkpoint(state, file_name):
|
313 |
+
torch.save(state, file_name)
|
314 |
+
|
315 |
+
|
316 |
+
if __name__ == '__main__':
|
317 |
+
# 超参数
|
318 |
+
parser = argparse.ArgumentParser(description='PyTorch Siamese Example')
|
319 |
+
parser.add_argument('--result_dir', default='output', type=str,
|
320 |
+
help='Directory to store results')
|
321 |
+
parser.add_argument('--epochs', type=int, default=10, metavar='N',
|
322 |
+
help='number of epochs to train (default: 10)')
|
323 |
+
|
324 |
+
parser.add_argument("--disable_save_ckp", default=False, action='store_true',
|
325 |
+
help="disable to save checkpoint frequently")
|
326 |
+
parser.add_argument('--ckp_freq', type=int, default=5, metavar='N',
|
327 |
+
help='Checkpoint Frequency (default: 1)')
|
328 |
+
parser.add_argument("--disable_save_best_ckp", default=False, action='store_true',
|
329 |
+
help="disable to save best checkpoint")
|
330 |
+
|
331 |
+
parser.add_argument('--train_log_step', type=int, default=500, metavar='M',
|
332 |
+
help='Number of iterations after which to log the loss')
|
333 |
+
parser.add_argument('--margin', type=float, default=1.0, metavar='M',
|
334 |
+
help='margin for triplet loss (default: 1.0)')
|
335 |
+
|
336 |
+
parser.add_argument('--batch_size', type=int, default=64, metavar='N',
|
337 |
+
help='input batch size for training (default: 64)')
|
338 |
+
parser.add_argument('--lr', type=float, default=0.0001, metavar='LR',
|
339 |
+
help='learning rate (default: 0.0001)')
|
340 |
+
parser.add_argument('--dataset', type=str, default='mnist', metavar='M',
|
341 |
+
help='Dataset (default: mnist)')
|
342 |
+
parser.add_argument('--picture_resize', type=int, default=200, metavar='M',
|
343 |
+
help='size of the picture to reset (default: 200)')
|
344 |
+
parser.add_argument('--num_train_samples', type=int, default=50000, metavar='M',
|
345 |
+
help='number of training samples (default: 50000)')
|
346 |
+
parser.add_argument('--num_test_samples', type=int, default=10000, metavar='M',
|
347 |
+
help='number of test samples (default: 10000)')
|
348 |
+
|
349 |
+
global args, device
|
350 |
+
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
351 |
+
args = parser.parse_args()
|
352 |
+
|
353 |
+
main()
|
形近字.txt
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
一乙二
|