snsd0805/Distributed-Training-Example
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: training log

Browse Source
This commit is contained in:
Ting-Jun Wang 2024-05-29 04:25:15 +08:00
parent 17949bd1a6
commit bdaa1d4846
Signed by: snsd0805
GPG Key ID: 48D331A3D6160354
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
train/trainer.py
View File

@ -31,6 +31,7 @@ class Trainer():
self.criterion = criterion
def train(self, epoch_num):
print("Start traininig...")
for epoch in range(epoch_num):
self.model.train()
train_loss_sum = 0
@ -52,7 +53,7 @@ class Trainer():
self.optimizer.step()
self.optimizer.zero_grad()
print(f"[DEVICE {self.global_rank}] EPOCH {epoch} loss={train_loss_sum/len(self.loader)} acc={(train_correct_sum/train_item_counter).item()}")
print(f"[RANK {self.global_rank}] EPOCH {epoch} loss={train_loss_sum/len(self.loader)} acc={(train_correct_sum/train_item_counter).item()}")
def save(self, model_path):
torch.save(self.model.state_dict(), model_path)