Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,545 Bytes
4c954ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
"""
This file contains some useful functions for train / val.
"""
import os
import numpy as np
import torch
import random
from scalelsd.ssl.models.detector import ScaleLSD
################
## HDF5 utils ##
################
def parse_h5_data(h5_data):
""" Parse h5 dataset. """
output_data = {}
for key in h5_data.keys():
output_data[key] = np.array(h5_data[key])
return output_data
def fix_seeds(random_seed):
random.seed(random_seed)
np.random.seed(random_seed)
os.environ['PYTHONHASHSEED'] = str(random_seed)
torch.manual_seed(random_seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(random_seed)
torch.cuda.manual_seed_all(random_seed) # if use multi-GPU
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# torch.backends.cudnn.allow_tf32 = args.tf32
# torch.backends.cuda.matmul.allow_tf32 = args.tf32
# torch.backends.cudnn.deterministic = args.dtm
# os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
# torch.use_deterministic_algorithms(True)
def load_scalelsd_model(ckpt_path, device='cuda'):
"""load model"""
use_layer_scale = False if 'v1' in ckpt_path else True
model = ScaleLSD(gray_scale=True, use_layer_scale=use_layer_scale)
model = model.eval().to(device)
state_dict = torch.load(ckpt_path, map_location='cpu')
try:
model.load_state_dict(state_dict['model_state'])
except:
model.load_state_dict(state_dict)
return model
|