Ajout de l'option pour sauvegarder/charger un model

.gitignore

 .idea/
 .venv/
 __pycache__/
-*_old.py
+*_old.py
\ No newline at end of file
+model.pt
\ No newline at end of file

main.py

+import os
 import torch
-import torch.onnx
+# import torch.onnx
 import torch.nn as nn
 
 import data
@@ -8,13 +9,23 @@ import musicgenerator as mg
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 if __name__ == '__main__':
+    ###############################################################################
+    # Paramètres
+    ###############################################################################
     learning_rate = 20
+    batch_size = 16
+    model_path = None  # os.path.normpath('./model.pt')
+    sequence_length = 32
+    epochs = 100
+    nb_log_epoch = 5
+    nb_words = 100
+    temperature = 1.0
+
     print("device :", device)
 
     ###############################################################################
     # Prépare les données
     ###############################################################################
-    batch_size = 16
     corpus = data.Corpus('./')
     train_data = data.batchify(corpus.train, batch_size, device)
     val_data = data.batchify(corpus.valid, batch_size, device)
@@ -23,32 +34,33 @@ if __name__ == '__main__':
     # Construit le modèle
     ###############################################################################
     ntokens = len(corpus.dictionary)
-    model = mg.MusicGenerator(vocab_size=ntokens, dim_model=512, num_head=8).to(device)
-    criterion = nn.NLLLoss()
 
-    ###############################################################################
-    # Entraîne le modèle
-    ###############################################################################
-    sequence_length = 32
-    nb_log_epoch = 5
-    epochs = 100
-    mg.train(
-        model=model,
-        criterion=criterion,
-        ntokens=ntokens,
-        train_data=train_data,
-        val_data=val_data,
-        sequence_length=sequence_length,
-        lr=learning_rate,
-        epochs=epochs,
-        log_interval=(len(train_data) // sequence_length) // nb_log_epoch
-    )
+    if model_path:
+        print("Load model from", model_path)
+        with open(model_path, 'rb') as f:
+            model = torch.load(f, map_location=device)
+    else:
+        model = mg.MusicGenerator(vocab_size=ntokens, dim_model=512, num_head=8).to(device)
+        criterion = nn.NLLLoss()
+
+        ###############################################################################
+        # Entraîne le modèle
+        ###############################################################################
+        mg.train(
+            model=model,
+            criterion=criterion,
+            ntokens=ntokens,
+            train_data=train_data,
+            val_data=val_data,
+            sequence_length=sequence_length,
+            lr=learning_rate,
+            epochs=epochs,
+            log_interval=(len(train_data) // sequence_length) // nb_log_epoch
+        )
 
     ###############################################################################
     # Génère une mélodie
     ###############################################################################
-    nb_words = 100
-    temperature = 1.0
     model.eval()
     input = torch.randint(ntokens, (1, 1), dtype=torch.long).to(device)
     with torch.no_grad():