NotImplementedError() when finetuning remora model

[palakela]

Hello,

I am trying to finetune this model to call 6mA (dna_r10.4.1_e8.2_5khz_400bps_sup_v4.2.0_6ma_v2.pt), but independently from how many layers I try to freeze (I've also tried with 0) I run into this error. Any idea on which could be the issue?

Traceback (most recent call last):
  File "./bin/remora", line 8, in <module>
    sys.exit(run())
  File "./lib/python3.10/site-packages/remora/main.py", line 71, in run
    cmd_func(args)
  File "./lib/python3.10/site-packages/remora/parsers.py", line 857, in run_model_train
    train_model(
  File "./lib/python3.10/site-packages/remora/train_model.py", line 180, in train_model
    ckpt, model = model_util.continue_from_checkpoint(
  File "./lib/python3.10/site-packages/remora/model_util.py", line 247, in continue_from_checkpoint
    if ckpt["state_dict"] is None:
  File "./lib/python3.10/site-packages/torch/jit/_script.py", line 862, in __getitem__
    return self.forward_magic_method("__getitem__", idx)
  File "./lib/python3.10/site-packages/torch/jit/_script.py", line 855, in forward_magic_method
    raise NotImplementedError()
NotImplementedError

For reference, I am using remora v. 3.1.0 and this is the command:

remora model train \
	${wd}/data/prepData/train_dataset.jsn \
	--model ${wd}/data/ONT/ConvLSTM_w_ref.py \
	--finetune-path ${wd}/dna_r10.4.1_e8.2_5khz_400bps_sup_v4.2.0_6ma_v2.pt \
	--device 0 \
	--chunk-context 50 50 \
	--output-path ${wd}/data/models/train_results

[marcus1487]

The resume feature requires a checkpoint model not a torch script model (optimized for inference, and lacking training status). You can recreate the checkpoint with the following snippet. I will try to add this to the core API to make this a bit simpler.

model, model_metadata = model_util._raw_load_torchscript_model(model_path)
model_metadata["epoch"] = 0
state_dict = model.state_dict()
if "total_ops" in state_dict.keys():
    state_dict.pop("total_ops", None)
if "total_params" in state_dict.keys():
    state_dict.pop("total_params", None)
model_metadata["state_dict"] = state_dict
model_metadata["opt"] = None
torch.save(
    model_metadata,
    out_path,
)

[palakela]

Thanks for clarifying!

I've managed to produce the checkpoint file to use in the remora training command. Yet, I am running into this other error now. Any idea from where the issue is coming from?

******************** WARNING [11:35:38.198:MainProcess:MainThread:train_model.py:196]: Size mismatch between pretrained model and selected size. Using pretrained model size. ********************
Traceback (most recent call last):
  File "./bin/remora", line 8, in <module>
    sys.exit(run())
             ^^^^^
  File "./lib/python3.12/site-packages/remora/main.py", line 71, in run
    cmd_func(args)
  File "./lib/python3.12/site-packages/remora/parsers.py", line 857, in run_model_train
    train_model(
  File "./lib/python3.12/site-packages/remora/train_model.py", line 202, in train_model
    raise RemoraError(
remora.RemoraError: The chunk context of the pre-trained model and the dataset do not match.

Reference command line:

remora model train \
	${wd}/data/prepData/train_dataset.jsn \
	--model ${wd}/data/ONT/ConvLSTM_w_ref.py \
	--finetune-path ${wd}/model_checkpoint.pth \
	--device 0 \
	--chunk-context 50 50 \
	--output-path ${wd}/data/models/train_results

[marcus1487]

Training from a checkpoint file requires that the same data input size be used. Setting the --chunk-context to the same value as the pre-trained model should resolve this issue. We will look into setting these parameters automatically in from the pre-trained model in the future. The best way to check this value is to either load the model using the python API or export the .pt file to a dorado model using the remora model export command and viewing the metadata in the config.toml file produced. I will flag up making a remora model inspect command to print out this information a bit more easily in the future.

[palakela]

I have double checked, both the pre-trained model and the training dataset have the same chunk_context [100,100]. changing this parameter in the remora model train commant, does not solve the issue.

[marcus1487]

Could you post the exact command and error message here to help resolve the issue?

[palakela]

here the full command line I used:

remora dataset prepare \
	--output-path ${wd}data/prepData/MOCK_6mA \
	--refine-kmer-level-table ${wd}data/ONT/9mer_levels_v1.txt \
	--refine-rough-rescale \
	--motif A 0 \
	--mod-base-control \
	--max-chunks-per-read 20 \
	--num-extract-alignment-workers 24 \
	--num-extract-chunks-workers 24 \
	--chunk-context 100 100 \
	--kmer-context-bases 4 4 \
	${wd}data/6mA_unmeth.pod5 \
	${wd}data/6mA_unmeth.pass.bam

remora dataset prepare \
	--output-path ${wd}data/prepData/MOD_6mA \
	--refine-kmer-level-table ${wd}data/ONT/9mer_levels_v1.txt \
	--refine-rough-rescale \
	--motif A 0 \
	--mod-base a 6mA \
	--max-chunks-per-read 20 \
	--num-extract-alignment-workers 24 \
	--num-extract-chunks-workers 24 \
	--chunk-context 100 100 \
	--kmer-context-bases 4 4 \
	${wd}data/7_6mA.pod5 \
	${wd}data/7_6mA.pass.bam

remora dataset make_config \
	${wd}data/prepData/train_dataset.jsn \
	${wd}data/prepData/MOCK_6mA \
	${wd}data/prepData/MOD_6mA \
	--dataset-weights 1 1 \
	--log-filename ${wd}data/prepData/train_dataset.log

python make_checkpoint.py \
	--model dna_r10.4.1_e8.2_5khz_400bps_sup_v4.2.0_6ma_v2.pt \
	--output ${wd}data/models/checkpoint_model.pth

remora model train \
	${wd}data/prepData/train_dataset.jsn \
	--model ${wd}data/ONT/ConvLSTM_w_ref.py \
	--finetune-path ${wd}data/models/checkpoint_model.pth \
	--freeze-num-layers 15 \
	--device 0 \
	--chunk-context 100 100 \
	--output-path ${wd}data/models/train_results_freeze15 \
	--kmer-context-bases 4 4

make_checkpoint.py is the snipped you provided me, taking the .pt file in input and returning the model checkpoint as output

here the full error message arising after the last command:

[11:30:12.563] Seed selected is 1106960644
[11:30:12.637] Loading dataset from Remora dataset config
[11:30:15.316] Dataset summary:
                     size : 33,433,400
     modified_base_labels : True
                mod_bases : ['a']
           mod_long_names : ['6mA']
       kmer_context_bases : (4, 4)
            chunk_context : (100, 100)
                   motifs : [('A', 0)]
           reverse_signal : False
 chunk_extract_base_start : False
     chunk_extract_offset : 0
          sig_map_refiner : Loaded 9-mer table with 7 central position. Rough re-scaling will be executed.

[11:30:15.317] Loading model
[11:30:15.467] Model structure:
network(
  (sig_conv1): Conv1d(1, 4, kernel_size=(5,), stride=(1,))
  (sig_bn1): BatchNorm1d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (sig_conv2): Conv1d(4, 16, kernel_size=(5,), stride=(1,))
  (sig_bn2): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (sig_conv3): Conv1d(16, 64, kernel_size=(9,), stride=(3,))
  (sig_bn3): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (seq_conv1): Conv1d(36, 16, kernel_size=(5,), stride=(1,))
  (seq_bn1): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (seq_conv2): Conv1d(16, 64, kernel_size=(13,), stride=(3,))
  (seq_bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (merge_conv1): Conv1d(128, 64, kernel_size=(5,), stride=(1,))
  (merge_bn): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (lstm1): LSTM(64, 64)
  (lstm2): LSTM(64, 64)
  (fc): Linear(in_features=64, out_features=2, bias=True)
  (dropout): Dropout(p=0.3, inplace=False)
)
******************** WARNING [11:30:15.499:MainProcess:MainThread:train_model.py:196]: Size mismatch between pretrained model and selected size. Using pretrained model size. ********************
Traceback (most recent call last):
  File "./envs/remora_v3.1.0/bin/remora", line 8, in <module>
    sys.exit(run())
             ^^^^^
  File "./envs/remora_v3.1.0/lib/python3.12/site-packages/remora/main.py", line 71, in run
    cmd_func(args)
  File "./envs/remora_v3.1.0/lib/python3.12/site-packages/remora/parsers.py", line 857, in run_model_train
    train_model(
  File "./envs/remora_v3.1.0/lib/python3.12/site-packages/remora/train_model.py", line 202, in train_model
    raise RemoraError(
remora.RemoraError: The chunk context of the pre-trained model and the dataset do not match.