31 lines
947 B
Python
31 lines
947 B
Python
import tempfile
|
|
import os
|
|
import torch
|
|
|
|
def save_model_to_hdf5_group(model, f):
|
|
tempfd, tempfname = tempfile.mkstemp(prefix='tmp-torchmodel')
|
|
try:
|
|
os.close(tempfd)
|
|
torch.save(model, tempfname)
|
|
with open(tempfname, 'rb') as model_file:
|
|
model_data = model_file.read()
|
|
f.create_dataset('torchmodel', data=model_data)
|
|
finally:
|
|
os.unlink(tempfname)
|
|
|
|
def load_model_from_hdf5_group(f, map_location=None):
|
|
tempfd, tempfname = tempfile.mkstemp(prefix='tmp-torchmodel')
|
|
try:
|
|
os.close(tempfd)
|
|
with open(tempfname, 'wb') as model_file:
|
|
model_file.write(f['torchmodel'][()])
|
|
model = torch.load(tempfname, map_location=map_location)
|
|
return model
|
|
finally:
|
|
os.unlink(tempfname)
|
|
|
|
def set_gpu_memory_target(device, frac):
|
|
# This function is not required in PyTorch since it doesn't pre-allocate all GPU memory by default.
|
|
pass
|
|
|