r""" .. math:: \min_x || y - X * x ||^2
| 31 | |
| 32 | |
| 33 | class LinearRegressionTrain(BaseTrainProcess): |
| 34 | r""" |
| 35 | .. math:: |
| 36 | \min_x || y - X * x ||^2 |
| 37 | """ |
| 38 | |
| 39 | def __init__( |
| 40 | self, |
| 41 | dim_w: int, |
| 42 | dim_y: int, # dataset size |
| 43 | batch_size: int, |
| 44 | lr: float = 1e-3, |
| 45 | *args, |
| 46 | **kwargs, |
| 47 | ): |
| 48 | super().__init__(*args, **kwargs) |
| 49 | |
| 50 | self.dim_w = dim_w |
| 51 | self.dim_y = dim_y |
| 52 | self.batch_size = batch_size |
| 53 | self.lr = lr |
| 54 | self._register_var_to_save(["dim_w", "dim_y"]) |
| 55 | |
| 56 | self._construct_data() |
| 57 | |
| 58 | def _construct_data(self): |
| 59 | g = torch.Generator() |
| 60 | g.manual_seed(0) |
| 61 | self.gt_w = torch.randn(self.dim_w, generator=g) # (dim_w,) |
| 62 | |
| 63 | self.xs = torch.randn(self.dim_y, self.dim_w, generator=g) # (n, dim_w) |
| 64 | self.ys = self.xs @ self.gt_w # (n,) |
| 65 | |
| 66 | self.xs_valid = torch.randn(self.dim_y, self.dim_w, generator=g) # (n, dim_w) |
| 67 | self.ys_valid = self.xs_valid @ self.gt_w # (n,) |
| 68 | |
| 69 | self.xs_test = torch.randn(self.dim_y, self.dim_w, generator=g) # (n, dim_w) |
| 70 | self.ys_test = self.xs_test @ self.gt_w # (n,) |
| 71 | |
| 72 | @customer |
| 73 | def get_dataloaders(self): |
| 74 | dset = QuickDataset(xs=self.xs, ys=self.ys) |
| 75 | dset_valid = QuickDataset(xs=self.xs_valid, ys=self.ys_valid) |
| 76 | dset_test = QuickDataset(xs=self.xs_test, ys=self.ys_test) |
| 77 | |
| 78 | if self.process_info['distributed_run']: |
| 79 | self.train_sampler = torch.utils.data.distributed.DistributedSampler( |
| 80 | dataset=dset, |
| 81 | num_replicas=self.process_info['n_gpus'], |
| 82 | rank=self.process_info['rank'], |
| 83 | seed=0, |
| 84 | shuffle=True, |
| 85 | drop_last=False, |
| 86 | ) |
| 87 | self.valid_sampler = torch.utils.data.distributed.DistributedSampler( |
| 88 | dataset=dset_valid, |
| 89 | num_replicas=self.process_info['n_gpus'], |
| 90 | rank=self.process_info['rank'], |