WangA commited on
Commit
0da05b1
1 Parent(s): 8d184e5

Upload 6 files

Browse files
Files changed (7) hide show
  1. .gitattributes +1 -0
  2. characters.zip +3 -0
  3. dataset.py +80 -0
  4. prepare_data.py +100 -0
  5. simhei.ttf +3 -0
  6. train.py +353 -0
  7. 形近字.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
+ 一乙二