In the test, we will use linear regression to test BaseTrainProcess.
| 251 | |
| 252 | |
| 253 | class TestBaseTrainProcess(unittest.TestCase): |
| 254 | """ |
| 255 | In the test, we will use linear regression to test BaseTrainProcess. |
| 256 | """ |
| 257 | |
| 258 | @staticmethod |
| 259 | def _test( |
| 260 | rank=0, |
| 261 | dim_w=30, |
| 262 | num_samples=1000, |
| 263 | batch_size=32, |
| 264 | lr=1e-1, |
| 265 | num_epoch=100, |
| 266 | exp_tag='test', |
| 267 | n_gpus=0, |
| 268 | ): |
| 269 | |
| 270 | with tempfile.TemporaryDirectory() as output_dir: |
| 271 | with LinearRegressionTrain( |
| 272 | dim_w=dim_w, |
| 273 | dim_y=num_samples, |
| 274 | rank=rank, |
| 275 | n_gpus=n_gpus, |
| 276 | batch_size=batch_size, |
| 277 | lr=lr, |
| 278 | exp_tag=exp_tag, |
| 279 | config_filename=None, |
| 280 | output_dir=output_dir, |
| 281 | end_epoch=num_epoch, |
| 282 | exp_tag_first=True, |
| 283 | log_every_num_train_batch=100, |
| 284 | log_every_num_valid_batch=100, |
| 285 | log_every_num_test_batch=100, |
| 286 | visualize_every_num_train_batch=1000, |
| 287 | visualize_every_num_valid_batch=1000, |
| 288 | visualize_every_num_test_batch=1000, |
| 289 | validate_every_num_epoch=10, |
| 290 | test_every_num_epoch=num_epoch, |
| 291 | use_bolt=False, |
| 292 | use_torchrun=False, |
| 293 | ) as trainer: |
| 294 | trainer.run() |
| 295 | gt_w = trainer.gt_w |
| 296 | |
| 297 | # load the trained model |
| 298 | if rank == 0: |
| 299 | filename = os.path.join( |
| 300 | output_dir, exp_tag, 'checkpoint', f'epoch{num_epoch - 1}.pth') |
| 301 | with LinearRegressionTrain( |
| 302 | dim_w=dim_w, |
| 303 | dim_y=num_samples, |
| 304 | batch_size=batch_size, |
| 305 | lr=lr, |
| 306 | exp_tag='test', |
| 307 | config_filename=None, |
| 308 | output_dir=output_dir, |
| 309 | ) as trainer: |
| 310 | trainer.construct_models() |
nothing calls this directly
no outgoing calls
no test coverage detected