File size: 2,011 Bytes
24a1704 762f6ed | 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 58 59 60 61 62 63 64 65 66 67 68 | #Author: Lart Pang (https://github.com/lartpang)
import logging
import os
import random
import numpy as np
import torch
from torch import nn
from torch.backends import cudnn
LOGGER = logging.getLogger("main")
def customized_worker_init_fn(worker_id):
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
def set_seed_for_lib(seed):
random.seed(seed)
np.random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def initialize_seed_cudnn(seed, deterministic):
assert isinstance(deterministic, bool) and isinstance(seed, int)
if seed >= 0:
LOGGER.info(f"We will use a fixed seed {seed}")
else:
seed = np.random.randint(2**32)
LOGGER.info(f"We will use a random seed {seed}")
set_seed_for_lib(seed)
if not deterministic:
LOGGER.info("We will use `torch.backends.cudnn.benchmark`")
else:
LOGGER.info("We will not use `torch.backends.cudnn.benchmark`")
cudnn.enabled = True
cudnn.benchmark = not deterministic
cudnn.deterministic = deterministic
def to_device(data, device="cuda"):
if isinstance(data, (tuple, list)):
return [to_device(item, device) for item in data]
elif isinstance(data, dict):
return {name: to_device(item, device) for name, item in data.items()}
elif isinstance(data, torch.Tensor):
return data.to(device=device, non_blocking=True)
else:
raise TypeError(f"Unsupported type {type(data)}. Only support Tensor or tuple/list/dict containing Tensors.")
def frozen_bn_stats(model, freeze_affine=False):
"""
Set all the bn layers to eval mode.
Args:
model (model): model to set bn layers to eval mode.
"""
for m in model.modules():
if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
m.eval()
if freeze_affine:
m.requires_grad_(False)
|