<jupyter_start><jupyter_text>Training code<jupyter_code>n_dataset = len(train_images)
train_data_size = 200 if debug else int(n_dataset * 0.85)
valid_data_size = 100 if debug else int(n_dataset - train_data_size)

perm = np.random.RandomState(777).permutation(n_dataset)
print('perm', perm)

train_dataset = BengaliAIDataset(
    train_images, train_labels, transform=train_transform,
    indices=perm[:train_data_size])

valid_dataset = BengaliAIDataset(
    train_images, train_labels, transform=Transform(affine=False, crop=False, size=(224, 224)),
    indices=perm[train_data_size:train_data_size+valid_data_size])

print('train_dataset', len(train_dataset), 'valid_dataset', len(valid_dataset))

'''
n_dataset = len(train_images)

train_data_size, valid_data_size= train_test_split(
np.arange(n_dataset),
test_size=0.2,
shuffle=True,
stratify=train_labels)

train_dataset = BengaliAIDataset(
    train_images, train_labels, transform=train_transform,
    indices=train_data_size)

valid_dataset = BengaliAIDataset(
    train_images, train_labels, transform=Transform(policy=False, affine=False, crop=False, size=(137, 236)),
    indices=valid_data_size)
print('train_dataset', len(train_dataset), 'valid_dataset', len(valid_dataset))
'''
class WrappedModel(nn.Module):
    def __init__(self, module):
        super(WrappedModel, self).__init__()
        self.module = module # that I actually define.
    def forward(self, x):
        return self.module(x)
# --- Model ---
device = torch.device(device)
n_grapheme = 168
n_vowel = 11
n_consonant = 7
n_root = 1295
n_total = n_grapheme + n_vowel + n_consonant + n_root
print('n_total', n_total)

predictor = PretrainedCNN(in_channels=3, out_dim=n_total, model_name=model_name, pretrained='imagenet')
#predictor = WrappedModel(predictor)
#load_checkpoint = torch.load('checkpoint.pt')
#predictor.load_state_dict(load_checkpoint['state_dict'])
#predictor = torch.load('checkpoint.pt')
print('predictor', type(predictor))

if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  predictor = nn.DataParallel(predictor)

classifier = BengaliClassifier(predictor).to(device)
classifier_val = BengaliClassifierVal(predictor).to(device)
def save_json(filepath, params):
    with open(filepath, 'w') as f:
        json.dump(params, f, indent=4)
class DictOutputTransform:
    def __init__(self, key, index=0):
        self.key = key
        self.index = index

    def __call__(self, x):
        if self.index >= 0:
            x = x[self.index]
        return x[self.key]


def create_trainer(classifier, optimizer, device):
    classifier.to(device)

    def update_fn(engine, batch):
        classifier.train()
        optimizer.zero_grad()
        x, y = [elem.to(device) for elem in batch]
        x, y_a, y_b, lam = cutmix(x, y, alpha)
        loss, metrics, pred_y = classifier(x, y, y_a, y_b, lam)
        loss.backward()
        optimizer.step()
        return metrics, pred_y, y
    trainer = Engine(update_fn)

    for key in classifier.metrics_keys:
        Average(output_transform=DictOutputTransform(key)).attach(trainer, key)
    return trainer


def create_evaluator(classifier, device):
    classifier.to(device)

    def update_fn(engine, batch):
        classifier.eval()
        with torch.no_grad():
            x, y = [elem.to(device) for elem in batch]
            _, metrics, pred_y = classifier(x, y)
            return metrics, pred_y, y
    evaluator = Engine(update_fn)

    for key in classifier.metrics_keys:
        Average(output_transform=DictOutputTransform(key)).attach(evaluator, key)
    return evaluator


class LogReport:
    def __init__(self, evaluator=None, dirpath=None, logger=None):
        self.evaluator = evaluator
        self.dirpath = str(dirpath) if dirpath is not None else None
        self.logger = logger or getLogger(__name__)

        self.reported_dict = {}  # To handle additional parameter to monitor
        self.history = []
        self.start_time = perf_counter()

    def report(self, key, value):
        self.reported_dict[key] = value

    def __call__(self, engine):
        elapsed_time = perf_counter() - self.start_time
        elem = {'epoch': engine.state.epoch,
                'iteration': engine.state.iteration}
        elem.update({f'train/{key}': value
                     for key, value in engine.state.metrics.items()})
        if self.evaluator is not None:
            elem.update({f'valid/{key}': value
                         for key, value in self.evaluator.state.metrics.items()})
        elem.update(self.reported_dict)
        elem['elapsed_time'] = elapsed_time
        self.history.append(elem)
        if self.dirpath:
            save_json(os.path.join(self.dirpath, 'log.json'), self.history)
            self.get_dataframe().to_csv(os.path.join(self.dirpath, 'log.csv'), index=False)

        # --- print ---
        msg = ''
        for key, value in elem.items():
            if key in ['iteration']:
                # skip printing some parameters...
                continue
            elif isinstance(value, int):
                msg += f'{key} {value: >6d} '
            else:
                msg += f'{key} {value: 8f} '
#         self.logger.warning(msg)
        print(msg)

        # --- Reset ---
        self.reported_dict = {}

    def get_dataframe(self):
        df = pd.DataFrame(self.history)
        return df


class SpeedCheckHandler:
    def __init__(self, iteration_interval=10, logger=None):
        self.iteration_interval = iteration_interval
        self.logger = logger or getLogger(__name__)
        self.prev_time = perf_counter()

    def __call__(self, engine: Engine):
        if engine.state.iteration % self.iteration_interval == 0:
            cur_time = perf_counter()
            spd = self.iteration_interval / (cur_time - self.prev_time)
            self.logger.warning(f'{spd} iter/sec')
            # reset
            self.prev_time = cur_time

    def attach(self, engine: Engine):
        engine.add_event_handler(Events.ITERATION_COMPLETED, self)


class ModelSnapshotHandler:
    def __init__(self, model, filepath='model/model_{count:03}.pt',
                 interval=1, logger=None):
        self.model = model
        self.filepath: str = str(filepath)
        self.interval = interval
        self.logger = logger or getLogger(__name__)
        self.count = 0

    def __call__(self, engine: Engine):
        self.count += 1
        if self.count % self.interval == 0:
            filepath = self.filepath.format(count=self.count)
            torch.save(self.model.state_dict(), filepath)
            checkpoint = {
            'epoch': engine.state.epoch,
            'state_dict': predictor.state_dict(),
            'optimizer': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict()    
            }
            torch.save(checkpoint, 'checkpoint.pt')
            # self.logger.warning(f'save model to {filepath}...')

class EpochMetric(Metric):
    """Class for metrics that should be computed on the entire output history of a model.
    Model's output and targets are restricted to be of shape `(batch_size, n_classes)`. Output
    datatype should be `float32`. Target datatype should be `long`.

    .. warning::

        Current implementation stores all input data (output and target) in as tensors before computing a metric.
        This can potentially lead to a memory error if the input data is larger than available RAM.


    - `update` must receive output of the form `(y_pred, y)`.

    If target shape is `(batch_size, n_classes)` and `n_classes > 1` than it should be binary: e.g. `[[0, 1, 0, 1], ]`.

    Args:
        compute_fn (callable): a callable with the signature (`torch.tensor`, `torch.tensor`) takes as the input
            `predictions` and `targets` and returns a scalar.
        output_transform (callable, optional): a callable that is used to transform the
            :class:`~ignite.engine.Engine`'s `process_function`'s output into the
            form expected by the metric. This can be useful if, for example, you have a multi-output model and
            you want to compute the metric with respect to one of the outputs.

    """

    def __init__(self, compute_fn, output_transform=lambda x: x):

        if not callable(compute_fn):
            raise TypeError("Argument compute_fn should be callable.")

        super(EpochMetric, self).__init__(output_transform=output_transform)
        self.compute_fn = compute_fn

    def reset(self):
        self._predictions = torch.tensor([], dtype=torch.float32)
        self._targets = torch.tensor([], dtype=torch.long)

    def update(self, output):
        y_pred, y = output
        self._predictions = torch.cat([self._predictions, y_pred], dim=0)
        self._targets = torch.cat([self._targets, y], dim=0)

        # Check once the signature and execution of compute_fn
        if self._predictions.shape == y_pred.shape:
            try:
                self.compute_fn(self._predictions, self._targets)
            except Exception as e:
                warnings.warn("Probably, there can be a problem with `compute_fn`:\n {}.".format(e),
                              RuntimeWarning)

    def compute(self):
        return self.compute_fn(self._predictions, self._targets)

def macro_recall(pred_y, y, n_grapheme=168, n_vowel=11, n_consonant=7, n_root=1295):
    pred_y = torch.split(pred_y, [n_grapheme, n_vowel, n_consonant, n_root], dim=1)
    pred_labels = [torch.argmax(py, dim=1).cpu().numpy() for py in pred_y]

    y = y.cpu().numpy()

    recall_grapheme = sklearn.metrics.recall_score(pred_labels[0], y[:, 0], average='macro')
    recall_vowel = sklearn.metrics.recall_score(pred_labels[1], y[:, 1], average='macro')
    recall_consonant = sklearn.metrics.recall_score(pred_labels[2], y[:, 2], average='macro')
    scores = [recall_grapheme, recall_vowel, recall_consonant]
    final_score = np.average(scores, weights=[2, 1, 1])
    print(f'recall: grapheme {recall_grapheme}, vowel {recall_vowel}, consonant {recall_consonant}, '
           f'total {final_score}, y {y.shape}')
    return final_score


def calc_macro_recall(solution, submission):
    scores = []
    for component in ['grapheme_root', 'consonant_diacritic', 'vowel_diacritic']:
        y_true_subset = solution[solution[component] == component]['target'].values
        y_pred_subset = submission[submission[component] == component]['target'].values
        scores.append(sklearn.metrics.recall_score(
            y_true_subset, y_pred_subset, average='macro'))
    final_score = np.average(scores, weights=[2, 1, 1])
    return final_score

# Store the best model
def default_score_fn(engine):
    score = engine.state.metrics['recall']
    return score

best_model_handler = ModelCheckpoint(dirname='best_models/',
                                     filename_prefix="best",
                                     n_saved=5,
                                     score_name="recall",
                                     score_function=default_score_fn,
                                     require_empty=False)
from ignite.engine import Engine, Events
from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=15)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=15)

optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)#0.001
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.7, patience=5, min_lr=1e-6)

#optimizer.load_state_dict(load_checkpoint['optimizer'])
#scheduler.load_state_dict(load_checkpoint['scheduler_state_dict'])
#epoch = load_checkpoint['epoch']

trainer = create_trainer(classifier, optimizer, device)
def output_transform(output):
    metric, pred_y, y = output
    return pred_y.cpu(), y.cpu()

EpochMetric(
    compute_fn=macro_recall,
    output_transform=output_transform
).attach(trainer, 'recall')

#pbar = ProgressBar()
#pbar.attach(trainer, metric_names='all')

evaluator = create_evaluator(classifier_val, device)
EpochMetric(
    compute_fn=macro_recall,
    output_transform=output_transform
).attach(evaluator, 'recall')

def run_evaluator(engine):
    evaluator.run(valid_loader)

def schedule_lr(engine):
    # metrics = evaluator.state.metrics
    metrics = engine.state.metrics
    avg_mae = metrics['loss']

    # --- update lr ---
    lr = scheduler.optimizer.param_groups[0]['lr']
    scheduler.step(avg_mae)
    log_report.report('lr', lr)

trainer.add_event_handler(Events.EPOCH_COMPLETED, run_evaluator)
trainer.add_event_handler(Events.EPOCH_COMPLETED, schedule_lr)
log_report = LogReport(evaluator, outdir)
#trainer.add_event_handler(Events.EPOCH_COMPLETED, log_report)
trainer.add_event_handler(
    Events.EPOCH_COMPLETED,
    #best_model_handler, {'model': predictor, })
    ModelSnapshotHandler(predictor))
    
#evaluator.add_event_handler(
#    Events.COMPLETED,
#    best_model_handler, {'model': predictor, })
    #ModelSnapshotHandler(predictor, filepath='model.pt'))


#to_save = {'model': predictor, 'optimizer': optimizer, 'lr_scheduler': scheduler, 'trainer': trainer}
#handler = Checkpoint(to_save, DiskSaver('tmp/models', create_dir=True), n_saved=5)
#trainer.add_event_handler(Events.EPOCH_COMPLETED, handler)
npt_logger.attach(trainer,
                  log_handler=OutputHandler(tag="training",
                                            output_transform=lambda loss: {'batchloss': loss},
                                            metric_names='all'),
                  event_name=Events.ITERATION_COMPLETED(every=100))

npt_logger.attach(evaluator,
                  log_handler=OutputHandler(tag="validation",
                                            metric_names='all',
                                            another_engine=trainer),
                  event_name=Events.EPOCH_COMPLETED)

npt_logger.attach(trainer,
                  log_handler=OptimizerParamsHandler(optimizer),
                  event_name=Events.ITERATION_COMPLETED(every=100))

trainer.run(train_loader, max_epochs=epochs)
train_history = log_report.get_dataframe()
train_history.to_csv(outdir / 'log.csv', index=False)

train_history
    
npt_logger.experiment.log_artifact(outdir / 'log.csv')
npt_logger.close()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11290434300755005, 'fraction_numerical': 0.007550045023204267, 'mean_word_length': 3.1381484666093438, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '31826', 'license': 'MIT', 'n_tokens_mistral': 4599, 'n_tokens_neox': 4467, 'n_words': 948, 'path': 'Bengali.Ai classification challenge/bengali-pytorch-mixup-train.ipynb', 'repo_name': 'yoviny/Kaggle-Competitions', 'size': '14437'}