コード例 #1
0
    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')
コード例 #2
0
    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")
コード例 #3
0
    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')