Let us train a GRBM and see how it performs
(args)
| 119 | |
| 120 | |
| 121 | def train_model(args): |
| 122 | """Let us train a GRBM and see how it performs""" |
| 123 | pid = os.getpid() |
| 124 | # Load config |
| 125 | with open(f'config/{args.dataset}.json') as json_file: |
| 126 | config = json.load(json_file) |
| 127 | |
| 128 | config['exp_folder'] = f"exp/{config['dataset']}_{config['model']}_{pid}_inference={config['inference_method']}_H={config['hidden_size']}_B={config['batch_size']}_CD={config['CD_step']}" |
| 129 | |
| 130 | if not os.path.isdir(config['exp_folder']): |
| 131 | os.makedirs(config['exp_folder']) |
| 132 | |
| 133 | log_file = os.path.join(config['exp_folder'], f'log_exp_{pid}.txt') |
| 134 | logger = setup_logging('INFO', log_file) |
| 135 | logger.info('Writing log file to {}'.format(log_file)) |
| 136 | |
| 137 | with open(os.path.join(config['exp_folder'], f'config_{pid}.json'), |
| 138 | 'w') as outfile: |
| 139 | json.dump(config, outfile, indent=4) |
| 140 | |
| 141 | config['visible_size'] = config['height'] * \ |
| 142 | config['width'] * config['channel'] |
| 143 | train_set = create_dataset(config) |
| 144 | train_loader = torch.utils.data.DataLoader(train_set, |
| 145 | batch_size=config['batch_size'], |
| 146 | shuffle=True) |
| 147 | |
| 148 | model = GRBM(config['visible_size'], |
| 149 | config['hidden_size'], |
| 150 | CD_step=config['CD_step'], |
| 151 | CD_burnin=config['CD_burnin'], |
| 152 | init_var=config['init_var'], |
| 153 | inference_method=config['inference_method'], |
| 154 | Langevin_step=config['Langevin_step'], |
| 155 | Langevin_eta=config['Langevin_eta'], |
| 156 | is_anneal_Langevin=True, |
| 157 | Langevin_adjust_step=config['Langevin_adjust_step']) |
| 158 | |
| 159 | if config['cuda']: |
| 160 | model.cuda() |
| 161 | |
| 162 | param_wd, param_no_wd = [], [] |
| 163 | for xx, yy in model.named_parameters(): |
| 164 | if 'W' in xx: |
| 165 | param_wd += [yy] |
| 166 | else: |
| 167 | param_no_wd += [yy] |
| 168 | |
| 169 | optimizer = optim.SGD([{ |
| 170 | 'params': param_no_wd, |
| 171 | 'weight_decay': 0 |
| 172 | }, { |
| 173 | 'params': param_wd |
| 174 | }], |
| 175 | lr=config['lr'], |
| 176 | momentum=0.0, |
| 177 | weight_decay=config['wd']) |
| 178 |
no test coverage detected