|
|
|
|
|
|
|
|
import os |
|
|
import os |
|
|
import random |
|
|
import random |
|
|
import time |
|
|
import time |
|
|
|
|
|
from copy import deepcopy |
|
|
from pathlib import Path |
|
|
from pathlib import Path |
|
|
from threading import Thread |
|
|
from threading import Thread |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
ckpt = {'epoch': epoch, |
|
|
ckpt = {'epoch': epoch, |
|
|
'best_fitness': best_fitness, |
|
|
'best_fitness': best_fitness, |
|
|
'training_results': results_file.read_text(), |
|
|
'training_results': results_file.read_text(), |
|
|
'model': (model.module if is_parallel(model) else model).half(), |
|
|
|
|
|
'ema': (ema.ema.half(), ema.updates), |
|
|
|
|
|
|
|
|
'model': deepcopy(model.module if is_parallel(model) else model).half(), |
|
|
|
|
|
'ema': (deepcopy(ema.ema).half(), ema.updates), |
|
|
'optimizer': optimizer.state_dict(), |
|
|
'optimizer': optimizer.state_dict(), |
|
|
'wandb_id': wandb_run.id if wandb else None} |
|
|
'wandb_id': wandb_run.id if wandb else None} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
torch.save(ckpt, best) |
|
|
torch.save(ckpt, best) |
|
|
del ckpt |
|
|
del ckpt |
|
|
|
|
|
|
|
|
model.float(), ema.ema.float() |
|
|
|
|
|
|
|
|
|
|
|
# end epoch ---------------------------------------------------------------------------------------------------- |
|
|
# end epoch ---------------------------------------------------------------------------------------------------- |
|
|
# end training |
|
|
# end training |
|
|
|
|
|
|