diff --git a/models/FineTuning/finetunning-multi-label.py b/models/FineTuning/finetunning-multi-label.py
index 420d00d2ba85a735c9cc08c395a8419ba18ecec0..164f305539e01cf256150d19452ac2f14d4f443d 100644
--- a/models/FineTuning/finetunning-multi-label.py
+++ b/models/FineTuning/finetunning-multi-label.py
@@ -19,7 +19,8 @@ from variables.pubmed import NCDS, NCDS_MESH_TERM
DATASET_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../testModel/dataset"))
TMP_DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../dataSources/PubMed/tmp"))
-DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "./data"))
+SAVE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "./save"))
+MODEL_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "./model"))
LOGS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "./logs"))
# Get PMIDs for all the articles from my testing dataset
@@ -73,15 +74,21 @@ tokenized_dataset = tokenized_dataset.remove_columns(["text"])
tokenized_dataset.set_format("torch")
training_args = TrainingArguments(
- output_dir = DATA_DIR,
+ output_dir = SAVE_DIR,
num_train_epochs = 32,
- per_device_train_batch_size = 12,
+ per_device_train_batch_size = 9,
warmup_steps = 500,
weight_decay = 0.01,
logging_dir = LOGS_DIR,
logging_steps=10,
fp16=True,
remove_unused_columns=False,
+ save_strategy="epoch",
+ # evaluation_strategy="epoch",
+ save_total_limit=3,
+ # load_best_model_at_end=True,
+ # metric_for_best_model="f1",
+ # greater_is_better=True
)
model = BartForSequenceClassification.from_pretrained(
@@ -94,12 +101,17 @@ model = BartForSequenceClassification.from_pretrained(
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print(f"Using device: {device}")
+for i in range(torch.cuda.device_count()):
+ print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
+
if torch.cuda.device_count() >