fusion-bench 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

fusion_bench_config/dataset/image_classification/train/TALL16.yaml

@@ -0,0 +1,28 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  # eight tasks in the task arithmetic paper
+  - sun397
+  - stanford-cars
+  - resisc45
+  - eurosat
+  - svhn
+  - gtsrb
+  - mnist
+  - dtd
+  # additional 6 tasks in the TALL mask paper (TALL 14)
+  - oxford_flowers102
+  - pcam
+  - fer2013
+  - oxford-iiit-pet
+  - stl10
+  - cifar100
+  # additional 6 tasks in the TALL mask paper (TALL 20)
+  - cifar10
+  - food101
+  # - fashion_mnist
+  # - emnist_letters
+  # - kmnist
+  # - rendered-sst2
+

fusion_bench_config/dataset/image_classification/train/TALL18.yaml

@@ -0,0 +1,28 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  # eight tasks in the task arithmetic paper
+  - sun397
+  - stanford-cars
+  - resisc45
+  - eurosat
+  - svhn
+  - gtsrb
+  - mnist
+  - dtd
+  # additional 6 tasks in the TALL mask paper (TALL 14)
+  - oxford_flowers102
+  - pcam
+  - fer2013
+  - oxford-iiit-pet
+  - stl10
+  - cifar100
+  # additional 6 tasks in the TALL mask paper (TALL 20)
+  - cifar10
+  - food101
+  - fashion_mnist
+  - emnist_letters
+  # - kmnist
+  # - rendered-sst2
+

fusion_bench_config/fabric_model_fusion.yaml

@@ -9,11 +9,11 @@ defaults:
 _target_: fusion_bench.programs.FabricModelFusionProgram
 _recursive_: false
 fast_dev_run: false # Run a single batch of data to test the model or method
-# Run the script without actually running the experiment, use with `print_config=true`. 
+# Run the script without actually running the experiment, use with `print_config=true`.
 # You can also use `--cfg` or `-c` to show the configuration instead of running.
 dry_run: false
 print_config: true # Print the configuration to the console
 merged_model_save_path: null # path to save the merged model, use "{log_dir}" to refer to the logger directory, for example `merged_model_save_path=\{log_dir\}/merged_model`
 merged_model_save_kwargs: null
-report_save_path: null # path to save the result report
+report_save_path: "{log_dir}/program_report.json" # path to save the result report
 print_function_call: true # set to false if you don't want to print the details of instantiate calls

fusion_bench_config/method/fw_merging/fw_hard.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.method.FrankWolfeHardAlgorithm
+merge_fn: task_arithmetic
+max_iters: 10
+step_size: 0.1
+dataset_size: 100
+tasks: []
+init_weight: 
+loss_fn: cross_entropy
+scaling_factor: 0.3
+max_num_models: 100
+granularity: task

fusion_bench_config/method/fw_merging/fw_soft.yaml

@@ -0,0 +1,12 @@
+_target_: fusion_bench.method.FrankWolfeSoftAlgorithm
+init_weight:
+max_iters: 10
+merge_fn: 'adamerging'
+tasks:
+ada_iters: 500
+dataset_size: 100
+ada_coeff: 1e-8
+step_size: 0.1
+max_num_models: 100
+granularity: task
+ada_loss: entropy_loss

fusion_bench_config/method/moe_pruner/moe_pruner.yaml

@@ -0,0 +1,15 @@
+_target_: fusion_bench.method.moe_pruner.MoEPruner
+
+nsamples: 100
+seed: 42
+device: cuda
+max_seqlen: 2048
+# `prune_type` can be either `unstructured` or `semistructured`
+prune_type: unstructured
+# === options for unstructured pruning ===
+# `sparsity_ratio` is the ratio of weights to be pruned, 1 means all weights are pruned
+sparsity_ratio: 0.5
+# === options for semistructured pruning ===
+# 2:4 means 2 out of 4 weights are pruned
+n: 2
+m: 4

fusion_bench_config/method/pruning/llama_sparsegpt_pruning.yaml

@@ -0,0 +1,16 @@
+_target_: fusion_bench.method.SparseGPTPruningForLlama
+nsamples: 128
+seed: 0
+use_variant: false
+# `prune_type` can be either `unstructured` or `semistructured`
+prune_type: unstructured
+# device and dtype to compute the pruning mask
+device: cuda
+dtype: null
+# === options for unstructured pruning ===
+# `sparsity_ratio` is the ratio of weights to be pruned, 1 means all weights are pruned
+sparsity_ratio: 0.5
+# === options for semistructured pruning ===
+# 2:4 means 2 out of 4 weights are pruned
+n: 2
+m: 4

fusion_bench_config/method/randes/superposed_model_soup.yaml

@@ -0,0 +1,18 @@
+_target_: fusion_bench.method.randes.SuperposedModelSoupAlgorithm
+#* === base randes options ===
+mode: manual_absorption
+# weights for all mlp and attn layers
+target_layer:
+  - mlp_w
+  - attn_w
+random_seed: 42 # for random_binary_diagonal_matrix
+different_across_layers: True
+joint_matrix_mode: flatten_hstack
+rank: 1 # for columnwise svd
+random_components: False
+shift_layers: 0
+absorber: None
+debug: 0
+ms_mode: average
+verbose: 0 # level of verbosity
+dropout_rate: 1 # take the target layer per n target layers

fusion_bench_config/method/randes/superposed_task_arithmetic.yaml

@@ -0,0 +1,20 @@
+name: superposed_task_arithmetic
+#* === base randes options ===
+mode: random_binary_diagonal_matrix
+# weights for all mlp and attn layers
+target_layer: 
+  - mlp_w
+  - attn_w
+random_seed: 42 # for random_binary_diagonal_matrix
+different_across_layers: True
+joint_matrix_mode: flatten_hstack
+rank: 1 # for columnwise svd
+random_components: False
+shift_layers: 0
+debug: 0
+verbose: 0
+dropout_rate: 1
+#* === task arithmetic options ===
+scaling_factor: 0.5
+# path to save/load the model
+model_path: null

fusion_bench_config/method/randes/superposed_task_arithmetic_lora.yaml

@@ -0,0 +1,20 @@
+_target_: fusion_bench.method.randes.SuperposedTaskArithmeticLoRAAlgorithm
+#* === base randes options ===
+mode: random_binary_diagonal_matrix
+# weights for all mlp and attn layers
+target_layer:
+  - mlp_w
+  - attn_w
+random_seed: 42 # for random_binary_diagonal_matrix
+different_across_layers: True
+joint_matrix_mode: flatten_hstack
+rank: 1 # for columnwise svd
+random_components: False
+shift_layers: 0
+debug: 0
+verbose: 0
+dropout_rate: 1
+#* === task arithmetic options ===
+scaling_factor: 0.5
+# path to save/load the model
+model_path: null

fusion_bench_config/method/sparselo_pruning/llama_iterative_sparselo.yaml

@@ -1,10 +1,11 @@
-_target_: fusion_bench.method.losparse.sparselo.IterativeSparseLoForLlama
+_target_: fusion_bench.method.sparselo.sparselo.IterativeSparseLoForLlama
 _recursive_: false
 nsamples: 128
 seed: 0
 rank: 128
 num_iterations: 10
 variant: wanda
+use_reference_model: false
 # `prune_type` can be either `unstructured` or `semistructured`
 prune_type: unstructured
 # device and dtype to compute the pruning mask

fusion_bench_config/method/sparselo_pruning/llama_pcp_sparselo.yaml

@@ -1,4 +1,4 @@
-_target_: fusion_bench.method.losparse.sparselo.PCPSparseLoForLlama
+_target_: fusion_bench.method.sparselo.sparselo.PCPSparseLoForLlama
 _recursive_: false
 nsamples: 128
 seed: 0

fusion_bench_config/method/sparselo_pruning/llama_sparselo.yaml

@@ -1,4 +1,4 @@
-_target_: fusion_bench.method.losparse.sparselo.SparseLoForLlama
+_target_: fusion_bench.method.sparselo.sparselo.SparseLoForLlama
 _recursive_: false
 nsamples: 128
 seed: 0

fusion_bench_config/method/tall_mask/task_arithmetic.yaml

@@ -0,0 +1,4 @@
+_target_: fusion_bench.method.tall_mask.TallMaskTaskArithmeticAlgorithm
+tall_mask_lambda: 0.6
+debug: 0
+verbose: 0

fusion_bench_config/method/task_singular_vector/TaskSingularVectorMerging.yaml

@@ -1,7 +1,8 @@
 _target_: fusion_bench.method.TaskSingularVectorMerging
-remove_keys: null
+exclude_keys: null
 # alpha is a float or a list of floats
 # example:
 # alpha: 1
 # alpha: [1, 0.5, 0.25]
 alpha: 1
+return_single_task_models: false

fusion_bench_config/model/clip-vit/clip-vit-base-patch32_TALL10.yaml

@@ -0,0 +1,29 @@
+# The 20 tasks used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  # pre-trained model
+  - clip-vit-base-patch32
+  # eight tasks in the task arithmetic paper
+  - clip-vit-base-patch32_sun397
+  - clip-vit-base-patch32_stanford-cars
+  - clip-vit-base-patch32_resisc45
+  - clip-vit-base-patch32_eurosat
+  - clip-vit-base-patch32_svhn
+  - clip-vit-base-patch32_gtsrb
+  - clip-vit-base-patch32_mnist
+  - clip-vit-base-patch32_dtd
+  # additional 6 tasks in the TALL mask paper (TALL 14)
+  - clip-vit-base-patch32_oxford_flowers102
+  - clip-vit-base-patch32_pcam
+  # - clip-vit-base-patch32_fer2013
+  # - clip-vit-base-patch32_oxford-iiit-pet
+  # - clip-vit-base-patch32_stl10
+  # - clip-vit-base-patch32_cifar100
+  # additional 6 tasks in the TALL mask paper (TALL 20)
+  # - clip-vit-base-patch32_cifar10
+  # - clip-vit-base-patch32_food101
+  # - clip-vit-base-patch32_fashion_mnist
+  # - clip-vit-base-patch32_emnist_letters
+  # - clip-vit-base-patch32_kmnist
+  # - clip-vit-base-patch32_rendered-sst2

fusion_bench_config/model/clip-vit/clip-vit-base-patch32_TALL12.yaml

@@ -0,0 +1,29 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  # pre-trained model
+  - clip-vit-base-patch32
+  # eight tasks in the task arithmetic paper
+  - clip-vit-base-patch32_sun397
+  - clip-vit-base-patch32_stanford-cars
+  - clip-vit-base-patch32_resisc45
+  - clip-vit-base-patch32_eurosat
+  - clip-vit-base-patch32_svhn
+  - clip-vit-base-patch32_gtsrb
+  - clip-vit-base-patch32_mnist
+  - clip-vit-base-patch32_dtd
+  # additional 6 tasks in the TALL mask paper (TALL 14)
+  - clip-vit-base-patch32_oxford_flowers102
+  - clip-vit-base-patch32_pcam
+  - clip-vit-base-patch32_fer2013
+  - clip-vit-base-patch32_oxford-iiit-pet
+  # - clip-vit-base-patch32_stl10
+  # - clip-vit-base-patch32_cifar100
+  # additional 6 tasks in the TALL mask paper (TALL 20)
+  # - clip-vit-base-patch32_cifar10
+  # - clip-vit-base-patch32_food101
+  # - clip-vit-base-patch32_fashion_mnist
+  # - clip-vit-base-patch32_emnist_letters
+  # - clip-vit-base-patch32_kmnist
+  # - clip-vit-base-patch32_rendered-sst2

fusion_bench_config/model/clip-vit/clip-vit-base-patch32_TALL16.yaml

@@ -0,0 +1,29 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  # pre-trained model
+  - clip-vit-base-patch32
+  # eight tasks in the task arithmetic paper
+  - clip-vit-base-patch32_sun397
+  - clip-vit-base-patch32_stanford-cars
+  - clip-vit-base-patch32_resisc45
+  - clip-vit-base-patch32_eurosat
+  - clip-vit-base-patch32_svhn
+  - clip-vit-base-patch32_gtsrb
+  - clip-vit-base-patch32_mnist
+  - clip-vit-base-patch32_dtd
+  # additional 6 tasks in the TALL mask paper (TALL 14)
+  - clip-vit-base-patch32_oxford_flowers102
+  - clip-vit-base-patch32_pcam
+  - clip-vit-base-patch32_fer2013
+  - clip-vit-base-patch32_oxford-iiit-pet
+  - clip-vit-base-patch32_stl10
+  - clip-vit-base-patch32_cifar100
+  # additional 6 tasks in the TALL mask paper (TALL 20)
+  - clip-vit-base-patch32_cifar10
+  - clip-vit-base-patch32_food101
+  # - clip-vit-base-patch32_fashion_mnist
+  # - clip-vit-base-patch32_emnist_letters
+  # - clip-vit-base-patch32_kmnist
+  # - clip-vit-base-patch32_rendered-sst2

fusion_bench_config/model/clip-vit/clip-vit-base-patch32_TALL18.yaml

@@ -0,0 +1,29 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  # pre-trained model
+  - clip-vit-base-patch32
+  # eight tasks in the task arithmetic paper
+  - clip-vit-base-patch32_sun397
+  - clip-vit-base-patch32_stanford-cars
+  - clip-vit-base-patch32_resisc45
+  - clip-vit-base-patch32_eurosat
+  - clip-vit-base-patch32_svhn
+  - clip-vit-base-patch32_gtsrb
+  - clip-vit-base-patch32_mnist
+  - clip-vit-base-patch32_dtd
+  # additional 6 tasks in the TALL mask paper (TALL 14)
+  - clip-vit-base-patch32_oxford_flowers102
+  - clip-vit-base-patch32_pcam
+  - clip-vit-base-patch32_fer2013
+  - clip-vit-base-patch32_oxford-iiit-pet
+  - clip-vit-base-patch32_stl10
+  - clip-vit-base-patch32_cifar100
+  # additional 6 tasks in the TALL mask paper (TALL 20)
+  - clip-vit-base-patch32_cifar10
+  - clip-vit-base-patch32_food101
+  - clip-vit-base-patch32_fashion_mnist
+  - clip-vit-base-patch32_emnist_letters
+  # - clip-vit-base-patch32_kmnist
+  # - clip-vit-base-patch32_rendered-sst2

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_TALL10.yaml

@@ -0,0 +1,8 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  - CLIPVisionModelPool@: _template
+  - /model/clip-vit@models: clip-vit-base-patch32_TALL10
+  - /dataset/image_classification/train@train_datasets: TALL10
+  - /dataset/image_classification/test@test_datasets: TALL10

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_TALL12.yaml

@@ -0,0 +1,8 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  - CLIPVisionModelPool@: _template
+  - /model/clip-vit@models: clip-vit-base-patch32_TALL12
+  - /dataset/image_classification/train@train_datasets: TALL12
+  - /dataset/image_classification/test@test_datasets: TALL12

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_TALL16.yaml

@@ -0,0 +1,8 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  - CLIPVisionModelPool@: _template
+  - /model/clip-vit@models: clip-vit-base-patch32_TALL16
+  - /dataset/image_classification/train@train_datasets: TALL16
+  - /dataset/image_classification/test@test_datasets: TALL16

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_TALL18.yaml

@@ -0,0 +1,8 @@
+# The 20 task used in the paper: 
+# Wang et al. Localizing Task Information for Improved Model Merging and Compression
+# http://arxiv.org/abs/2405.07813
+defaults:
+  - CLIPVisionModelPool@: _template
+  - /model/clip-vit@models: clip-vit-base-patch32_TALL18
+  - /dataset/image_classification/train@train_datasets: TALL18
+  - /dataset/image_classification/test@test_datasets: TALL18

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_cars_and_dtd.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - /model/clip-vit@models:
+      - clip-vit-base-patch32
+      - clip-vit-base-patch32_stanford-cars
+      - clip-vit-base-patch32_dtd
+  - /dataset/image_classification/train@train_datasets:
+      - stanford-cars
+      - dtd
+  - /dataset/image_classification/test@test_datasets:
+      - stanford-cars
+      - dtd
+_target_: fusion_bench.modelpool.CLIPVisionModelPool
+_recursive_: False
+processor:
+  _target_: transformers.CLIPProcessor.from_pretrained
+  pretrained_model_name_or_path: openai/clip-vit-base-patch32

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_sun397_and_cars.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - /model/clip-vit@models:
+      - clip-vit-base-patch32
+      - clip-vit-base-patch32_sun397
+      - clip-vit-base-patch32_stanford-cars
+  - /dataset/image_classification/train@train_datasets:
+      - sun397
+      - stanford-cars
+  - /dataset/image_classification/test@test_datasets:
+      - sun397
+      - stanford-cars
+_target_: fusion_bench.modelpool.CLIPVisionModelPool
+_recursive_: False
+processor:
+  _target_: transformers.CLIPProcessor.from_pretrained
+  pretrained_model_name_or_path: openai/clip-vit-base-patch32

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_sun397_and_dtd.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - /model/clip-vit@models:
+      - clip-vit-base-patch32
+      - clip-vit-base-patch32_sun397
+      - clip-vit-base-patch32_dtd
+  - /dataset/image_classification/train@train_datasets:
+      - sun397
+      - dtd
+  - /dataset/image_classification/test@test_datasets:
+      - sun397
+      - dtd
+_target_: fusion_bench.modelpool.CLIPVisionModelPool
+_recursive_: False
+processor:
+  _target_: transformers.CLIPProcessor.from_pretrained
+  pretrained_model_name_or_path: openai/clip-vit-base-patch32

fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_sun397_cars_and_dtd.yaml

@@ -0,0 +1,19 @@
+defaults:
+  - /model/clip-vit@models:
+      - clip-vit-base-patch32
+      - clip-vit-base-patch32_sun397
+      - clip-vit-base-patch32_stanford-cars
+      - clip-vit-base-patch32_dtd
+  - /dataset/image_classification/train@train_datasets:
+      - sun397
+      - stanford-cars
+      - dtd
+  - /dataset/image_classification/test@test_datasets:
+      - sun397
+      - stanford-cars
+      - dtd
+_target_: fusion_bench.modelpool.CLIPVisionModelPool
+_recursive_: False
+processor:
+  _target_: transformers.CLIPProcessor.from_pretrained
+  pretrained_model_name_or_path: openai/clip-vit-base-patch32

fusion_bench_config/modelpool/CausalLMPool/deepseek-v2-lite.yaml

@@ -0,0 +1,15 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+
+pretrained_model_name_or_path: deepseek-ai/DeepSeek-V2-Lite
+
+models:
+  _pretrained_:
+    _target_: fusion_bench.models.modeling_deepseek_v2.DeepseekV2ForCausalLM.from_pretrained
+    pretrained_model_name_or_path: ${...pretrained_model_name_or_path}
+    torch_dtype: bfloat16
+    device_map: auto
+    trust_remote_code: true
+
+tokenizer:
+  _target_: transformers.AutoTokenizer.from_pretrained
+  pretrained_model_name_or_path: ${..pretrained_model_name_or_path}

fusion_bench_config/modelpool/CausalLMPool/mixtral-8x7b.yaml

@@ -0,0 +1,14 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+
+pretrained_model_name_or_path: mistralai/Mixtral-8x7B-v0.1
+
+models:
+  _pretrained_:
+    _target_: transformers.AutoModelForCausalLM.from_pretrained
+    pretrained_model_name_or_path: ${...pretrained_model_name_or_path}
+    torch_dtype: bfloat16
+    device_map: auto
+
+tokenizer:
+  _target_: transformers.AutoTokenizer.from_pretrained
+  pretrained_model_name_or_path: ${..pretrained_model_name_or_path}

fusion_bench_config/modelpool/SeqenceClassificationModelPool/roberta-base_glue.yaml

@@ -0,0 +1,69 @@
+defaults:
+  - Seq2SeqLMPool@: _template
+  - /model/roberta@models:
+      - roberta_base
+      - roberta_glue-cola
+      - roberta_glue-mnli
+      - roberta_glue-mrpc
+      - roberta_glue-qnli
+      - roberta_glue-qqp
+      - roberta_glue-rte
+      - roberta_glue-sst2
+      - roberta_glue-stsb
+# _target_: fusion_bench.modelpool.SequenceClassificationModelPool
+# _recursive_: false
+
+_dataset_loader: fusion_bench.tasks.flan_t5_text_generation.glue_load_dataset.load_glue_dataset
+test_datasets:
+  glue-cola:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: cola
+    tokenizer: ${...tokenizer}
+    split: validation
+  glue-mnli:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: mnli
+    tokenizer: ${...tokenizer}
+    split: validation_matched
+  glue-mrpc:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: mrpc
+    tokenizer: ${...tokenizer}
+    split: validation
+  glue-qnli:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: qnli
+    tokenizer: ${...tokenizer}
+    split: validation
+  glue-qqp:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: qqp
+    tokenizer: ${...tokenizer}
+    split: validation
+  glue-rte:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: rte
+    tokenizer: ${...tokenizer}
+    split: validation
+  glue-sst2:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: sst2
+    tokenizer: ${...tokenizer}
+    split: validation
+  glue-stsb:
+    _target_: ${..._dataset_loader}
+    _recursive_: false
+    name: stsb
+    tokenizer: ${...tokenizer}
+    split: validation
+
+tokenizer:
+  _target_: transformers.AutoTokenizer.from_pretrained
+  pretrained_model_name_or_path: roberta-base

fusion_bench_config/taskpool/LMEvalHarnessTaskPool/lm_eval.yaml

@@ -1,5 +1,4 @@
 _target_: fusion_bench.taskpool.LMEvalHarnessTaskPool
-
 tasks:
   - truthfulqa
 batch_size: 1