(choices)
| 11 | import time |
| 12 | |
| 13 | def sample_configs(choices): |
| 14 | |
| 15 | config = {} |
| 16 | dimensions = ['mlp_ratio', 'num_heads'] |
| 17 | depth = random.choice(choices['depth']) |
| 18 | for dimension in dimensions: |
| 19 | config[dimension] = [random.choice(choices[dimension]) for _ in range(depth)] |
| 20 | |
| 21 | config['embed_dim'] = [random.choice(choices['embed_dim'])]*depth |
| 22 | |
| 23 | config['layer_num'] = depth |
| 24 | return config |
| 25 | |
| 26 | def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module, |
| 27 | data_loader: Iterable, optimizer: torch.optim.Optimizer, |
no test coverage detected