main.py

"""
后台训练：python main.py -u > output.log 2>&1 &
"""
from torch.optim.adam import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.nn.modules.loss import CrossEntropyLoss

from data import CaptchaDataset, DataLoader
from model import CaptchaModel


# 超参数
learning_rate = 0.0001
batch_size = 256
epoch = 1
# 初始化模型
model = CaptchaModel()
# 优化器
optimizer = Adam(model.nn.parameters(), lr=learning_rate)
# 损失函数
criterion = CrossEntropyLoss()
# 学习率调度器
scheduler = ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
# 训练集
train_dataset = CaptchaDataset()
# 验证集
validation_dataset = CaptchaDataset(kind='validation')

val_loss = 0
for i in range(epoch):
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    train_accuracy = model.train(train_dataloader, optimizer, criterion, accuracy=True)
    val_loss, val_accuracy = model.evaluate(validation_dataset, criterion)
    scheduler.step(val_loss)
    model.save(
        f'./trained-models/captcha{i + 1}.pth',
        learning_rate=optimizer.param_groups[0]['lr'],
        epoch=epoch,
        batch_size=batch_size
    )
    print(f'当前迭代次数：{i + 1} 验证集损失：{val_loss} 训练集准确率：{train_accuracy} 验证集准确率：{val_accuracy}')

test_dataset = CaptchaDataset(kind='test')

loss, accuracy = model.evaluate(test_dataset, criterion)
print(loss, accuracy)

# import os
# import cv2
# from model import CaptchaModel
#
#
# model = CaptchaModel('cpu')
# model.load('./captcha175.pth')
#
# rights = 0
# for filename in os.listdir('./dataset/test'):
#     image = cv2.imread('./dataset/test/' + filename)
#     predict = model.predict(image)
#     true = filename.split('.')[0]
#     if predict == true:
#         rights += 1
#     else:
#         print(predict, true)
# print(rights / 2000)

# import cv2
# from model import CaptchaModel
#
#
# model = CaptchaModel('cpu')
# model.load('./trained-models/captcha-v2.pth')
#
# image = cv2.imread('./dataset/test/0R87gL.png')
# predict = model.predict(image)
# print(predict)