| 146 | |
| 147 | |
| 148 | def init_model(opt): |
| 149 | |
| 150 | if opt.grayscale: |
| 151 | out_features = 1 |
| 152 | else: |
| 153 | out_features = 3 |
| 154 | |
| 155 | if opt.model == 'mlp': |
| 156 | |
| 157 | if opt.multiscale: |
| 158 | m = modules.MultiscaleCoordinateNet |
| 159 | else: |
| 160 | m = modules.CoordinateNet |
| 161 | |
| 162 | model = m(nl=opt.activation, |
| 163 | in_features=2, |
| 164 | out_features=out_features, |
| 165 | hidden_features=opt.hidden_features, |
| 166 | num_hidden_layers=opt.hidden_layers, |
| 167 | w0=opt.w0, |
| 168 | pe_scale=opt.pe_scale, |
| 169 | no_pe=opt.no_pe, |
| 170 | integrated_pe=opt.ipe) |
| 171 | |
| 172 | elif opt.model == 'mfn': |
| 173 | |
| 174 | if opt.multiscale: |
| 175 | m = modules.MultiscaleBACON |
| 176 | else: |
| 177 | m = modules.BACON |
| 178 | |
| 179 | input_scales = [1/8, 1/8, 1/4, 1/4, 1/4] |
| 180 | output_layers = [1, 2, 4] |
| 181 | |
| 182 | model = m(2, opt.hidden_features, out_size=out_features, |
| 183 | hidden_layers=opt.hidden_layers, |
| 184 | bias=True, |
| 185 | frequency=(opt.res, opt.res), |
| 186 | quantization_interval=2*np.pi, |
| 187 | input_scales=input_scales, |
| 188 | output_layers=output_layers, |
| 189 | reuse_filters=False) |
| 190 | |
| 191 | else: |
| 192 | raise ValueError('model must be mlp or mfn') |
| 193 | |
| 194 | model_parameters = filter(lambda p: p.requires_grad, model.parameters()) |
| 195 | params = sum([np.prod(p.size()) for p in model_parameters]) |
| 196 | print(f'Num. Parameters: {params}') |
| 197 | model.cuda() |
| 198 | |
| 199 | return model |
| 200 | |
| 201 | |
| 202 | def init_loss(opt, trn_dataset, val_dataset): |