def __init__(
self,
betas=(0.5, 0.999),
criterion_class='torch.nn.MSELoss',
init_weights=True,
lr=0.001,
nn_class='fnet.nn_modules.fnet_nn_3d.Net',
nn_kwargs={},
scheduler=None,
weight_decay=0,
gpu_ids=-1,
):
self.betas = betas
self.criterion = str_to_class(criterion_class)()
self.gpu_ids = [gpu_ids] if isinstance(gpu_ids, int) else gpu_ids
self.init_weights = init_weights
self.lr = lr
self.nn_class = nn_class
self.nn_kwargs = nn_kwargs
self.scheduler = scheduler
self.weight_decay = weight_decay
self.count_iter = 0
self.device = (torch.device('cuda', self.gpu_ids[0])
if self.gpu_ids[0] >= 0 else torch.device('cpu'))
self.optimizer = None
self._init_model()
self.fnet_model_kwargs, self.fnet_model_posargs = get_args()
self.fnet_model_kwargs.pop('self')
def __init__(
self,
betas=(0.5, 0.999),
criterion_class="fnet.losses.WeightedMSE",
init_weights=True,
lr=0.001,
nn_class="fnet.nn_modules.fnet_nn_3d.Net",
nn_kwargs={},
scheduler=None,
weight_decay=0,
gpu_ids=-1,
):
self.betas = betas
self.criterion = str_to_object(criterion_class)()
self.gpu_ids = [gpu_ids] if isinstance(gpu_ids, int) else gpu_ids
self.init_weights = init_weights
self.lr = lr
self.nn_class = nn_class
self.nn_kwargs = nn_kwargs
self.scheduler = scheduler
self.weight_decay = weight_decay
self.count_iter = 0
self.device = (
torch.device("cuda", self.gpu_ids[0])
if self.gpu_ids[0] >= 0
else torch.device("cpu")
)
self.optimizer = None
self._init_model()
self.fnet_model_kwargs, self.fnet_model_posargs = get_args()
self.fnet_model_kwargs.pop("self")
def __init__(
self,
betas=(0.5, 0.999),
criterion_class='torch.nn.MSELoss',
init_weights=True,
lr=0.001,
nn_class='fnet.nn_modules.fnet_nn_3d.Net',
nn_kwargs={},
nn_module=None,
scheduler=None,
weight_decay=0,
gpu_ids=-1,
):
self.betas = betas
self.criterion = str_to_class(criterion_class)()
self.gpu_ids = [gpu_ids] if isinstance(gpu_ids, int) else gpu_ids
self.init_weights = init_weights
self.lr = lr
self.nn_class = nn_class
self.nn_kwargs = nn_kwargs
self.scheduler = scheduler
self.weight_decay = weight_decay
# *** Legacy support ***
# self.nn_module might be specified in legacy saves.
# If so, override self.nn_class
if nn_module is not None:
self.nn_class = nn_module + '.Net'
self.count_iter = 0
self.device = (
torch.device('cuda', self.gpu_ids[0])
if self.gpu_ids[0] >= 0
else torch.device('cpu')
)
self._init_model()
self.fnet_model_kwargs, self.fnet_model_posargs = get_args()
self.fnet_model_kwargs.pop('self')
