Update train.py

train.py

@@ -96,7 +96,7 @@ class BaseDset(object):
         self.__test_keys = []
 
     def load(self, base_path):
-        """加载数据集，将类别和路径存储"""
+        """加载训练和测试数据集，将类别和路径存储"""
         self.__base_path = base_path
         train_dir = os.path.join(self.__base_path, 'train')
         test_dir = os.path.join(self.__base_path, 'test')
@@ -107,17 +107,14 @@ class BaseDset(object):
         self.__test_keys = []
 
         for class_id in os.listdir(train_dir):
-            # 对于train_dir里的每个文件夹名字 classi
             class_dir = os.path.join(train_dir, class_id)
-            # 为其在训练集合中创建一个文件夹
-            # 在类别集合中，即train_keys中添加类别classi
             self.__train_set[class_id] = []
             self.__train_keys.append(class_id)
-            # 对于每个类别内的数据，将其路径添加到集合中
+
             for img_name in os.listdir(class_dir):
                 img_path = os.path.join(class_dir, img_name)
                 self.__train_set[class_id].append(img_path)
-        # 同理对于测试集合也一样
+
         for class_id in os.listdir(test_dir):
             class_dir = os.path.join(test_dir, class_id)
             self.__test_set[class_id] = []
@@ -128,7 +125,7 @@ class BaseDset(object):
 
         return len(self.__train_keys), len(self.__test_keys)
 
-    # 获取三元组 ！！！
+    # 获取三元组
     def getTriplet(self, split='train'):
         # 默认选取训练集
         if split == 'train':
@@ -167,11 +164,11 @@ def train(data, model, criterion, optimizer, epoch):
     total_loss = 0
     model.train()
     for batch_idx, img_triplet in enumerate(data):
-        # 提取数据
+        # 提取三元组数据
         anchor_img, pos_img, neg_img = img_triplet
         anchor_img, pos_img, neg_img = anchor_img.to(device), pos_img.to(device), neg_img.to(device)
         anchor_img, pos_img, neg_img = Variable(anchor_img), Variable(pos_img), Variable(neg_img)
-        # 分别获得三个编码
+        # 分别获得三个编码，表示原始样本、正样本、负样本
         E1, E2, E3 = model(anchor_img, pos_img, neg_img)
         # 计算二者之间的欧式距离
         dist_E1_E2 = F.pairwise_distance(E1, E2, 2)
@@ -180,14 +177,14 @@ def train(data, model, criterion, optimizer, epoch):
         target = torch.FloatTensor(dist_E1_E2.size()).fill_(-1)
         target = target.to(device)
         target = Variable(target)
-        # 大小如何？
+
         loss = criterion(dist_E1_E2, dist_E1_E3, target)
         total_loss += loss
 
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
-        # 打印一波损失
+        # 打印损失
         log_step = args.train_log_step
         if (batch_idx % log_step == 0) and (batch_idx != 0):
             print('Train Epoch: {} [{}/{}] \t Loss: {:.4f}'.format(epoch, batch_idx, len(data), total_loss / log_step))
@@ -225,7 +222,7 @@ def test(data, model, criterion):
                 accuracies[i] += batch_acc
         print('Test Loss: {}'.format(total_loss / len(data)))
         for i in range(len(accuracies)):
-            # 0%等价于准确率其余是更严格的指标
+            # 0%等价于准确率，其余是更严格的指标
             print(
                 'Test Accuracy with diff = {}% of margin: {:.4f}'.format(acc_threshes[i] * 100,
                                                                          accuracies[i] / len(data)))