| 36 | self.writer = SummaryWriter(log_dir="saved/BGRL_dataset({})".format(args.name)) |
| 37 | |
| 38 | def _init(self): |
| 39 | args = self._args |
| 40 | os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device) |
| 41 | self._device = f'cuda:{args.device}' if torch.cuda.is_available() else "cpu" |
| 42 | # self._dataset = data.Dataset(root=args.root, name=args.name)[0] |
| 43 | self._dataset = data.get_data(args.name) |
| 44 | print(f"Data: {self._dataset}") |
| 45 | hidden_layers = [int(dim) for dim in args.layers] |
| 46 | layers = [self._dataset.x.shape[1]] + hidden_layers |
| 47 | self._model = models.BGRL(layer_config=layers, pred_hid=args.pred_hid, dropout=args.dropout, epochs=args.epochs).to(self._device) |
| 48 | print(self._model) |
| 49 | |
| 50 | self._optimizer = optim.AdamW(params=self._model.parameters(), lr=args.lr, weight_decay=1e-5) |
| 51 | |
| 52 | # learning rate |
| 53 | def lr_scheduler(epoch): |
| 54 | if epoch <= args.warmup_epochs: |
| 55 | return epoch / args.warmup_epochs |
| 56 | else: |
| 57 | return (1 + np.cos((epoch - args.warmup_epochs) * np.pi / (self._args.epochs - args.warmup_epochs))) * 0.5 |
| 58 | # lr_scheduler = lambda epoch: epoch / args.warmup_epochs if epoch <= args.warmup_epochs \ |
| 59 | # else ( 1 + np.cos((epoch - args.warmup_epochs) * np.pi / (self._args.epochs - args.warmup_epochs))) * 0.5 |
| 60 | |
| 61 | self._scheduler = optim.lr_scheduler.LambdaLR(self._optimizer, lr_lambda=lr_scheduler) |
| 62 | |
| 63 | def train(self): |
| 64 | # get initial test results |