添加fine_tuning，以及训练时可修改参数

captcha_train.py

@@ -18,7 +18,7 @@ def main():
     optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
 
     # Train the Model
-    train_dataloader = my_dataset.get_train_data_loader()
+    train_dataloader = my_dataset.get_train_data_loader(batch_size=batch_size)
     for epoch in range(num_epochs):
         for i, (images, labels) in enumerate(train_dataloader):
             images = Variable(images)

fine_tuning.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import my_dataset
+from captcha_cnn_model import CNN
+
+learning_rate = 0.001
+batch_size = 10
+num_epochs = 15
+
+
+def fine_tuning():
+    cnn = CNN()
+    cnn.eval()
+    cnn.load_state_dict(torch.load('model.pkl'))
+    print("load cnn net.")
+    criterion = nn.MultiLabelSoftMarginLoss()
+    optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
+    train_data_loader = my_dataset.get_train_data_loader(batch_size)
+
+    for epoch in range(num_epochs):
+        for i, (images, labels) in enumerate(train_data_loader):
+            images = Variable(images)
+            labels = Variable(labels.float())
+            predict_labels = cnn(images)
+            # print(predict_labels.type)
+            # print(labels.type)
+            loss = criterion(predict_labels, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if (i + 1) % 10 == 0:
+                print("epoch:", epoch, "step:", i, "loss:", loss.item())
+            if (i + 1) % 100 == 0:
+                torch.save(cnn.state_dict(), "./model.pkl")  # current is model.pkl
+                print("save model")
+        print("epoch:", epoch, "step:", i, "loss:", loss.item())
+    torch.save(cnn.state_dict(), "./model.pkl")  # current is model.pkl
+    print("save last model")
+
+if __name__ == '__main__':
+    fine_tuning()
+

my_dataset.py

@@ -30,10 +30,10 @@ def __getitem__(self, idx):
     transforms.ToTensor(),
     # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
-def get_train_data_loader():
+def get_train_data_loader(batch_size=64):
 
     dataset = mydataset(captcha_setting.TRAIN_DATASET_PATH, transform=transform)
-    return DataLoader(dataset, batch_size=64, shuffle=True)
+    return DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
 def get_test_data_loader():
     dataset = mydataset(captcha_setting.TEST_DATASET_PATH, transform=transform)