| import torch |
| import os |
| import torch.nn.functional as F |
|
|
| def save_checkpoint(model, filelocation, save_parallel = True): |
| if save_parallel: |
| torch.save(model.module.state_dict(), filelocation) |
| else: |
| torch.save(model.state_dict(), filelocation) |
|
|
| def load_Checkpoint(fileLocation,model, load_cpu=False): |
| if load_cpu: |
| model.load_state_dict(torch.load(fileLocation,map_location=lambda storage, loc: storage)) |
| else: |
| model.load_state_dict(torch.load(fileLocation)) |
| return model |
|
|
| def writeLog(logList, filename): |
| with open(filename, 'w') as outfile: |
| outfile.write("\n".join(logList)) |
|
|
|
|
| def kl_loss(mu, logvar): |
| return -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).mean() |
|
|
|
|
|
|
