Release training script

Former-commit-id: 784d960221ee0a4799cfc44ead65ce515294c716

model/LISA.py

@@ -6,15 +6,15 @@ import torch.nn.functional as F
 from peft import LoraConfig, get_peft_model
 from transformers import BitsAndBytesConfig, CLIPVisionModel
 
-from transformers import CLIPVisionModel, BitsAndBytesConfig
-from .llava.model.llava import LlavaLlamaForCausalLM
-from .segment_anything import build_sam_vit_h
 from utils.utils import (
     DEFAULT_IM_END_TOKEN,
     DEFAULT_IM_START_TOKEN,
     DEFAULT_IMAGE_PATCH_TOKEN,
 )
 
+from .llava.model.llava import LlavaLlamaForCausalLM
+from .segment_anything import build_sam_vit_h
+
 
 def dice_loss(
     inputs: torch.Tensor,

train_ds.py

@@ -63,7 +63,7 @@ def parse_args(args):
     parser.add_argument("--dataset_dir", default="./dataset", type=str)
     parser.add_argument("--log_base_dir", default="./runs", type=str)
     parser.add_argument("--exp_name", default="lisa", type=str)
-    parser.add_argument("--epochs", default=20, type=int)
+    parser.add_argument("--epochs", default=10, type=int)
     parser.add_argument("--steps_per_epoch", default=500, type=int)
     parser.add_argument(
         "--batch_size", default=2, type=int, help="batch size per device per step"