| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from model_cnn import PlainCNN |
| from dataloader2 import SAnime2XDataset |
|
|
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
| n_layer = 12 |
| lr_peak = 2.5e-4 |
| n_steps = 10000 |
|
|
| batch_size = 8 |
| patch_size = 96 |
| zoom = 2 |
|
|
| ds = SAnime2XDataset('/home/bctnry/s/sanime2x-2k/') |
|
|
| model = PlainCNN(n_layer) |
| model = model.to(device) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=lr_peak) |
| scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 2500, gamma=0.9) |
|
|
| PSNR_EVAL_IDX = torch.randint(ds.val_size, (2,)) |
| def val_psnr(model): |
| calc = [] |
| idx = PSNR_EVAL_IDX |
| images = ds.get_val_images(idx) |
| se_images = ds.get_val_shrink_enlarged_images(idx, zoom=zoom) |
| model.eval() |
| with torch.no_grad(): |
| for k, img in enumerate(images): |
| se_img = se_images[k] |
| h, w = img.shape[1], img.shape[2] |
| p_h, p_w = h // patch_size, w // patch_size |
| |
| |
| for j in range(0, p_h): |
| for i in range(0, p_w): |
| cropped = img[:, |
| j*patch_size:(j+1)*patch_size, |
| i*patch_size:(i+1)*patch_size] |
| cropped = cropped.reshape(1, *cropped.shape) |
| cropped = cropped.to(device) |
| se_cropped = se_img[:, |
| j*patch_size:(j+1)*patch_size, |
| i*patch_size:(i+1)*patch_size] |
| se_cropped = se_cropped.reshape(1, *se_cropped.shape) |
| se_cropped = se_cropped.to(device) / 255 |
| res = model(se_cropped) |
| cropped = cropped / 255 |
| |
| loss = F.mse_loss(res, cropped) |
| calc.append(loss.item()) |
| total_loss = sum(calc) / len(calc) |
| psnr = -10*torch.log10(torch.Tensor([total_loss])).item() |
| return psnr |
|
|
| losses = [] |
| avg_span = 50 |
| last_psnr = 0 |
| for i in range(n_steps): |
| model.train() |
| |
| enlarged, cropped = ds.get_train_random_crop(n=batch_size, zoom=zoom) |
| enlarged = enlarged.to(device) |
| cropped = cropped.to(device) |
| calculated = model(enlarged) |
| |
| loss = F.l1_loss(calculated, cropped) |
|
|
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
| scheduler.step() |
| |
| losses.append(loss.item()) |
| if len(losses) > avg_span: |
| losses = losses[1:] |
|
|
| avgloss = sum(losses)/len(losses) |
| print(f'Step {i+1}: loss={loss.item():.4f} avgloss={avgloss:.4f} last_psnr={last_psnr:.4f} dB') |
| |
| |
| if i % 500 == 499: |
| last_psnr = val_psnr(model) |
| if i % 100 == 99: |
| torch.save({ |
| 'model': model.state_dict(), |
| 'optim': optimizer.state_dict(), |
| }, 'model_cnn.pt') |
| |
| torch.save({ |
| 'model': model.state_dict(), |
| 'optim': optimizer.state_dict(), |
| }, 'model_cnn.pt') |
| |
|
|
|
|
