dataset preparation
(opt)
| 25 | |
| 26 | |
| 27 | def train(opt): |
| 28 | """ dataset preparation """ |
| 29 | if not opt.data_filtering_off: |
| 30 | print('Filtering the images containing characters which are not in opt.character') |
| 31 | print('Filtering the images whose label is longer than opt.batch_max_length') |
| 32 | # see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130 |
| 33 | |
| 34 | opt.select_data = opt.select_data.split('-') |
| 35 | opt.batch_ratio = opt.batch_ratio.split('-') |
| 36 | train_dataset = Batch_Balanced_Dataset(opt) |
| 37 | |
| 38 | log = open(f'./saved_models/{opt.exp_name}/log_dataset.txt', 'a') |
| 39 | AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD) |
| 40 | valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt) |
| 41 | valid_loader = torch.utils.data.DataLoader( |
| 42 | valid_dataset, batch_size=opt.batch_size, |
| 43 | shuffle=True, # 'True' to check training progress with validation function. |
| 44 | num_workers=int(opt.workers), |
| 45 | collate_fn=AlignCollate_valid, pin_memory=True) |
| 46 | log.write(valid_dataset_log) |
| 47 | print('-' * 80) |
| 48 | log.write('-' * 80 + '\n') |
| 49 | log.close() |
| 50 | |
| 51 | """ model configuration """ |
| 52 | if 'CTC' in opt.Prediction: |
| 53 | if opt.baiduCTC: |
| 54 | converter = CTCLabelConverterForBaiduWarpctc(opt.character) |
| 55 | else: |
| 56 | converter = CTCLabelConverter(opt.character) |
| 57 | else: |
| 58 | converter = AttnLabelConverter(opt.character) |
| 59 | |
| 60 | opt.num_class = len(converter.character) |
| 61 | |
| 62 | if opt.rgb: |
| 63 | opt.input_channel = 3 |
| 64 | model = Model(opt) |
| 65 | print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel, |
| 66 | opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction, |
| 67 | opt.SequenceModeling, opt.Prediction) |
| 68 | |
| 69 | # weight initialization |
| 70 | # for name, param in model.named_parameters(): |
| 71 | # if 'localization_fc2' in name: |
| 72 | # print(f'Skip {name} as it is already initialized') |
| 73 | # continue |
| 74 | # try: |
| 75 | # if 'bias' in name: |
| 76 | # init.constant_(param, 0.0) |
| 77 | # elif 'weight' in name: |
| 78 | # init.kaiming_normal_(param) |
| 79 | # except Exception as e: # for batchnorm. |
| 80 | # if 'weight' in name: |
| 81 | # param.data.fill_(1) |
| 82 | # continue |
| 83 | |
| 84 | # data parallel for multi-GPU |
no test coverage detected