os.environ["CUDA_VISIBLE_DEVICES"] = "0" log_filename = os.path.join('log/', 'loss_acc-' + config.saved_model_prefix + '.log') if not os.path.exists('debug_files'): os.mkdir('debug_files') if not os.path.exists(config.saved_model_dir): os.mkdir(config.saved_model_dir) if config.use_log and not os.path.exists('log'): os.mkdir('log') if config.use_log and os.path.exists(log_filename): os.remove(log_filename) if config.experiment is None: config.experiment = 'expr' if not os.path.exists(config.experiment): os.mkdir(config.experiment)
设置和检查文件和目录的一些条件。具体的功能包括：
1. 设置环境变量`CUDA_VISIBLE_DEVICES`为`0`，表示使用第一个GPU设备。
2. 定义日志文件名`log_filename`，用于保存模型的损失和准确率。
3. 检查是否存在`debug_files`目录，如果不存在则创建。
4. 检查是否存在`config.saved_model_dir`目录，如果不存在则创建。该目录用于保存模型的参数和状态。
5. 检查是否启用日志功能并且是否存在`log`目录，如果启用并且目录不存在则创建。
6. 如果启用日志功能并且`log_filename`已经存在，则删除该文件。
7. 检查是否指定了实验名称`config.experiment`，如果没有则设置为默认名称`expr`。
8. 检查是否存在`config.experiment`目录，如果不存在则创建。
train_dataset = mydataset.MyDataset(info_filename=config.train_infofile) assert train_dataset if optimizer_type == 'rms': config.manualSeed = random.randint(1, 10000) # fix seed print("Random Seed: ", config.manualSeed) random.seed(config.manualSeed) np.random.seed(config.manualSeed) torch.manual_seed(config.manualSeed) sampler = mydataset.randomSequentialSampler(train_dataset, config.batchSize) else: sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.batchSize, shuffle=True, sampler=sampler, num_workers=int(config.workers), collate_fn=mydataset.alignCollate(imgH=config.imgH, imgW=config.imgW, keep_ratio=config.keep_ratio)) test_dataset = mydataset.MyDataset( info_filename=config.val_infofile, transform=mydataset.resizeNormalize((config.imgW, config.imgH), is_test=True)) converter = utils.strLabelConverter(config.alphabet) criterion = CTCLoss(reduction='sum', zero_infinity=True)
设置训练过程中的数据集加载器、数据转换器、损失函数等。
1. 创建`mydataset.MyDataset`的实例`train_dataset`，并传入训练数据的信息文件路径`config.train_infofile`。
2. 使用`assert`语句检查`train_dataset`是否存在（非空）。
3. 如果优化器类型`optimizer_type`为`rms`，则执行以下操作：
- 生成一个随机种子`config.manualSeed`，范围为1到10000，用于固定随机数生成器的种子。
- 打印出随机种子的值`config.manualSeed`。
- 使用`random.seed()`、`np.random.seed()`和`torch.manual_seed()`设置相应的随机数种子。
- 使用`mydataset.randomSequentialSampler()`创建一个随机顺序的采样器`sampler`，该采样器用于对训练数据进行采样，每次迭代选择一个随机的样本。
4. 如果优化器类型不是`rms`，则`sampler`为`None`。
5. 创建`torch.utils.data.DataLoader`的实例`train_loader`，用于加载训练数据集。参数包括：
- `train_dataset`：训练数据集。
- `batch_size`：批大小，即每次加载的样本数量。
- `shuffle`：是否在每个epoch中打乱数据集。
- `sampler`：采样器，用于决定样本的顺序。
- `num_workers`：加载数据的线程数。
- `collate_fn`：用于对样本进行对齐和规范化处理的函数。
6. 创建`mydataset.MyDataset`的实例`test_dataset`，并传入验证数据的信息文件路径`config.val_infofile`，以及对图像进行尺寸调整和标准化的数据转换器`mydataset.resizeNormalize()`。
7. 创建一个字符标签转换器`converter`，用于将字符标签转换为模型可处理的张量形式。
8. 创建CTC损失函数`criterion`，设置`reduction='sum'`表示对每个样本的损失求和，`zero_infinity=True`表示遇到无穷大的损失时将其设置为0。
def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: m.weight.data.normal_(0.0, 0.02) elif classname.find('BatchNorm') != -1: m.weight.data.normal_(1.0, 0.02) m.bias.data.fill_(0) crnn = crnn.CRNN(config.imgH, config.nc, config.nclass, config.nh) if config.pretrained_model != '' and os.path.exists(config.pretrained_model): print('loading pretrained model from %s' % config.pretrained_model) crnn.load_state_dict(torch.load(config.pretrained_model)) else: crnn.apply(weights_init)
初始化模型权重并创建CRNN模型。
1. 定义了一个`weights_init`函数，用于对模型的权重进行初始化。该函数通过判断`m`的类名是否包含关键字`Conv`或`BatchNorm`来确定执行的初始化操作。对于`Conv`层，使用正态分布随机初始化权重；对于`BatchNorm`层，使用正态分布随机初始化权重，并将偏置项设置为0。
2. 创建一个CRNN模型的实例`crnn`，并传入图像的高度`config.imgH`、通道数`config.nc`、字符类别数`config.nclass`和隐藏层尺寸`config.nh`作为参数。
3. 如果指定了预训练模型文件路径`config.pretrained_model`并且该文件存在，则加载该预训练模型的权重到`crnn`模型中。
4. 如果没有指定预训练模型或者预训练模型文件不存在，则对`crnn`模型的权重进行初始化，调用`crnn.apply(weights_init)`函数。
device = torch.device('cpu') if config.cuda: crnn.cuda() # crnn = torch.nn.DataParallel(crnn, device_ids=range(opt.ngpu)) # image = image.cuda() device = torch.device('cuda:0') criterion = criterion.cuda()
将模型和损失函数移动到设备上进行计算。
1. 首先将默认设备设置为CPU，通过`device = torch.device('cpu')`实现。
2. 检查配置文件中的`cuda`标志是否为真，如果为真，则将模型和损失函数移动到CUDA设备上进行计算。
- 调用`crnn.cuda()`将模型移动到CUDA设备上。
- 创建一个`torch.device`对象表示CUDA设备，通过`device = torch.device('cuda:0')`指定使用第一个CUDA设备。
- 调用`criterion.cuda()`将损失函数移动到CUDA设备上。
loss_avg = utils.averager() nbs = config.batchSize * 8 lr_limit_max = 1e-3 if optimizer_type == 'adam' else 5e-2 lr_limit_min = 3e-4 if optimizer_type == 'adam' else 5e-4 Init_lr_fit = min(max(config.batchSize / nbs * Init_lr, lr_limit_min), lr_limit_max) Min_lr_fit = min(max(config.batchSize / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2) lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, Epoch)
初始化学习率相关的参数和学习率调度器。
1. 创建一个`utils.averager()`实例`loss_avg`，用于计算平均损失。
2. 计算每个epoch中的总迭代次数`nbs`，即批大小乘以8倍的批大小。
3. 根据优化器类型设置学习率的上下限值。如果优化器类型为'adam'，则将学习率上限`lr_limit_max`设置为1e-3，否则设置为5e-2；将学习率下限`lr_limit_min`设置为3e-4（'adam'优化器）或5e-4（其他优化器）。
4. 根据批大小和总迭代次数计算初始学习率`Init_lr_fit`和最小学习率`Min_lr_fit`，并对其进行限制在上下限范围内。
- 初始学习率`Init_lr_fit`计算方法：将配置文件中的初始学习率`Init_lr`乘以批大小除以总迭代次数，在上下限范围内取最大值。
- 最小学习率`Min_lr_fit`计算方法：将配置文件中的最小学习率`Min_lr`乘以批大小除以总迭代次数，在上下限范围内取最大值，并乘以1e-2。
5. 调用`get_lr_scheduler`函数生成学习率调度器`lr_scheduler_func`。调度器的参数包括：
- `lr_decay_type`：学习率衰减类型，用于确定调度器的类型。
- `Init_lr_fit`：初始学习率。
- `Min_lr_fit`：最小学习率。
- `Epoch`：当前的训练epoch数。
optimizer = { 'adam' : optim.Adam(crnn.parameters(), Init_lr_fit, betas = (momentum, 0.999)), 'sgd' : optim.SGD(crnn.parameters(), Init_lr_fit, momentum = momentum, nesterov=True), 'rms' : optim.RMSprop(crnn.parameters(), lr=Init_lr_fit) }[optimizer_type]
根据选择的优化器类型创建相应的优化器对象。
根据`optimizer_type`选择不同的优化器类型：
- 如果`optimizer_type`为'adam'，则创建一个Adam优化器对象`optim.Adam(crnn.parameters(), Init_lr_fit, betas=(momentum, 0.999))`。该优化器使用CRNN模型的参数作为优化的目标，并设置初始学习率为`Init_lr_fit`，动量系数为`momentum`，beta系数为（`momentum`，0.999）。
- 如果`optimizer_type`为'sgd'，则创建一个SGD优化器对象`optim.SGD(crnn.parameters(), Init_lr_fit, momentum=momentum, nesterov=True)`。该优化器使用CRNN模型的参数作为优化的目标，并设置初始学习率为`Init_lr_fit`，动量系数为`momentum`，使用Nesterov加速。
- 如果`optimizer_type`为'rms'，则创建一个RMSprop优化器对象`optim.RMSprop(crnn.parameters(), lr=Init_lr_fit)`。该优化器使用CRNN模型的参数作为优化的目标，并设置初始学习率为`Init_lr_fit`。
def val(net, dataset, criterion, max_iter=100): print('Start val') for p in net.parameters(): p.requires_grad = False num_correct, num_all = val_model(config.val_infofile, net, True, log_file='compare-' + config.saved_model_prefix + '.log') accuracy = num_correct / num_all print('ocr_acc: %f' % (accuracy)) if config.use_log: with open(log_filename, 'a') as f: f.write('ocr_acc:{}\n'.format(accuracy)) global best_acc if accuracy > best_acc: best_acc = accuracy torch.save(crnn.state_dict(), '{}/{}_{}_{}.pth'.format(config.saved_model_dir, config.saved_model_prefix, epoch, int(best_acc * 1000))) torch.save(crnn.state_dict(), '{}/{}.pth'.format(config.saved_model_dir, config.saved_model_prefix))
定义了一个验证函数`val`，用于在验证集上评估模型的性能。
函数参数：
- `net`：需要评估的模型。
- `dataset`：验证数据集。
- `criterion`：损失函数。
- `max_iter`：最大迭代次数，默认为100。
函数主要步骤：
1. 打印开始验证的提示信息。
2. 将模型中的所有参数的`requires_grad`属性设置为`False`，固定模型参数不进行梯度更新。
3. 调用`val_model`函数对模型进行验证，计算正确的样本数`num_correct`和总样本数`num_all`。同时将验证结果打印到日志文件中，文件名为'compare-' + config.saved_model_prefix + '.log'。
4. 计算准确率`accuracy`，即正确的样本数除以总样本数。
5. 打印准确率。
6. 如果配置中设置了使用日志文件`config.use_log`，则将准确率结果写入日志文件中。
7. 将当前准确率与最佳准确率`best_acc`进行比较，如果当前准确率大于最佳准确率，则更新最佳准确率，并保存模型参数。保存的模型参数文件名包括配置中的保存模型目录`config.saved_model_dir`、保存模型前缀`config.saved_model_prefix`、当前epoch和准确率乘以1000的值。
8. 不论是否更新最佳准确率，都保存一份模型参数文件，文件名为保存模型目录加上保存模型前缀。
def trainBatch(net, criterion, optimizer): data = next(train_iter) cpu_images, cpu_texts = data batch_size = cpu_images.size(0) image = cpu_images.to(device) text, length = converter.encode(cpu_texts) # utils.loadData(text, t) # utils.loadData(length, l) preds = net(image) # seqLength x batchSize x alphabet_size preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size)) # seqLength x batchSize cost = criterion(preds.log_softmax(2).cpu(), text, preds_size, length) / batch_size if torch.isnan(cost): print(batch_size, cpu_texts) else: net.zero_grad() cost.backward() optimizer.step() return cost
定义了一个训练一个batch的函数`trainBatch`，用于对模型进行一次batch的训练。
函数参数：
- `net`：需要训练的模型。
- `criterion`：损失函数。
- `optimizer`：优化器，用于更新模型的参数。
函数主要步骤：
1. 取出一个batch的训练数据，包括图像和对应的标签。
2. 将图像数据转移到设备上。
3. 将标签使用转码器进行编码，得到编码后的标签和标签长度。
4. 使用模型对图像进行前向传播，得到预测结果`preds`，其维度为(seqLength x batchSize x alphabet_size)。
5. 创建变量`preds_size`，用于指定预测结果的尺寸，将其设置为(seqLength x batchSize)。
6. 计算损失值`cost`，将预测结果进行log softmax操作，然后与编码后的标签、预测结果尺寸和标签长度一起传入损失函数计算损失值，并除以batch_size得到平均损失。
7. 如果损失值是NaN（无效值），则打印出batch的大小和对应的文本内容。
8. 否则，将模型的梯度清零，然后进行反向传播计算梯度，并使用优化器更新模型参数。
9. 返回损失值`cost`。
for epoch in range(config.niter): loss_avg.reset() str_lr = set_optimizer_lr(optimizer, lr_scheduler_func, epoch) print('epoch {}....lr {}'.format(epoch, str_lr)) train_iter = iter(train_loader) i = 0 n_batch = len(train_loader) while i < len(train_loader): for p in crnn.parameters(): p.requires_grad = True crnn.train() cost = trainBatch(crnn, criterion, optimizer) print('epoch: {} iter: {}/{} Train loss: {:.3f}'.format(epoch, i, n_batch, cost.item())) loss_avg.add(cost) loss_avg.add(cost) i += 1 print('Train loss: %f' % (loss_avg.val())) if config.use_log: with open(log_filename, 'a') as f: f.write('{}\n'.format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f'))) f.write('train_loss:{}\n'.format(loss_avg.val())) val(crnn, test_dataset, criterion)
训练循环，用于训练模型多个epoch，并在每个epoch结束后进行验证。
代码主要步骤：
1. 对损失计算器`loss_avg`进行重置，用于计算每个epoch的平均损失。
2. 调用`set_optimizer_lr`函数设置优化器的学习率，并将学习率的字符串表示赋值给`str_lr`变量。
3. 打印当前epoch和学习率。
4. 获得训练数据集的迭代器`train_iter`。
5. 初始化循环变量`i`为0，并计算总batch数`n_batch`。
6. 在每个batch内循环，直到处理完所有的batch：
- 将模型参数的`requires_grad`属性设置为`True`，允许模型参数进行梯度更新。
- 将模型设置为训练模式。
- 调用`trainBatch`函数对模型进行训练，并得到训练损失值`cost`。
- 打印当前epoch、当前batch索引和总batch数，以及训练损失值。
- 将训练损失值添加到损失计算器`loss_avg`中。
- 增加循环变量`i`的值。
7. 打印当前epoch的平均训练损失。
8. 如果配置中设置了使用日志文件`config.use_log`，则将当前时间和训练损失值写入日志文件中。
9. 调用`val`函数对模型在验证集上进行验证。