| 155 | |
| 156 | |
| 157 | def optimize(): |
| 158 | optimizer = torch.optim.Adam(params, lr=0.01) |
| 159 | sigmoid = torch.nn.Sigmoid() |
| 160 | for iteration in range(ITERATIONS): |
| 161 | optimizer.zero_grad() |
| 162 | |
| 163 | blobs = torch.concat( |
| 164 | ( |
| 165 | blob_positions, |
| 166 | torch.exp(blob_scales), |
| 167 | sigmoid(blob_colors), |
| 168 | blob_rotations, |
| 169 | ), |
| 170 | dim=1, |
| 171 | ) |
| 172 | image = model.forward(blobs) |
| 173 | loss = torch.nn.functional.l1_loss(image, target_cuda) |
| 174 | # loss = torch.nn.functional.mse_loss(image, target_cuda) |
| 175 | loss.backward() |
| 176 | optimizer.step() |
| 177 | |
| 178 | if iteration % 5 == 0: |
| 179 | print(f"iteration={iteration}, loss={loss.item()}") |
| 180 | render_image = image.detach() |
| 181 | # render_image = target_cuda.detach() |
| 182 | # render_image = torch.lerp(image.detach(), target_cuda.detach(), 0.5) |
| 183 | # render_image = torch.abs(image.detach() - target_cuda.detach()) |
| 184 | render_image = torch.pow(render_image, 2.2) |
| 185 | testbed.render_texture.from_numpy(render_image.cpu().numpy()) |
| 186 | testbed.frame() |
| 187 | if testbed.should_close: |
| 188 | break |
| 189 | |
| 190 | |
| 191 | # run optimization |