PyTorch による Noise2Noise の実装と動作確認（画像のノイズ除去）（PyTorch，Python 3.12 を使用）（Windows 上）

前準備

winget install -e --id Python.Python.3.12 --scope machine --silent --accept-source-agreements --accept-package-agreements --override "/quiet InstallAllUsers=1 PrependPath=1 AssociateFiles=1 InstallLauncherAllUsers=1"

python --version

PyTorch のインストール

NVIDIA GPU を使用する場合，NVIDIA CUDA ツールキットや cuDNN を個別にシステムへインストールする必要はない．PyTorch の GPU 版パッケージには，必要な CUDA 関連ライブラリが同梱されている．必要なのは，GPU に対応した最新の NVIDIA ドライバのみである．

管理者権限のコマンドプロンプトで，次のコマンドを実行する．

python -m pip install -U torch torchvision --index-url https://download.pytorch.org/whl/cu126
python -m pip install -U numpy pillow

【CUDA バージョンの選択について】 上記の cu126 は，CUDA 12.6 に対応した PyTorch パッケージを意味する．使用している GPU・ドライバに応じて，PyTorch 公式サイトのインストールページ（https://pytorch.org/get-started/locally/）で該当するコマンドを確認すること．GPU を使用しない場合は，--index-url のオプションを省略してよい．

インストール後，次のコマンドで GPU が認識されているか確認する．

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

True と表示されれば，GPU（CUDA）が利用可能である．

Noise2Noise モデルの実装

U-Net 構造を用いた，最小構成の Noise2Noise モデルを実装する．次のコードをテキストエディタで作成し，n2n_model.py という名前で保存する．

import torch
import torch.nn as nn

class N2NUNet(nn.Module):
    """Noise2Noise 用の簡易 U-Net"""
    def __init__(self, in_channels=3, out_channels=3):
        super().__init__()

        def conv_block(in_ch, out_ch):
            return nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1),
                nn.LeakyReLU(0.1, inplace=True),
                nn.Conv2d(out_ch, out_ch, 3, padding=1),
                nn.LeakyReLU(0.1, inplace=True),
            )

        self.enc1 = conv_block(in_channels, 48)
        self.pool1 = nn.MaxPool2d(2)
        self.enc2 = conv_block(48, 96)
        self.pool2 = nn.MaxPool2d(2)
        self.enc3 = conv_block(96, 96)

        self.up2 = nn.Upsample(scale_factor=2, mode="nearest")
        self.dec2 = conv_block(96 + 96, 96)
        self.up1 = nn.Upsample(scale_factor=2, mode="nearest")
        self.dec1 = conv_block(96 + 48, 64)

        self.out_conv = nn.Conv2d(64, out_channels, 3, padding=1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool1(e1))
        e3 = self.enc3(self.pool2(e2))

        d2 = self.up2(e3)
        d2 = self.dec2(torch.cat([d2, e2], dim=1))
        d1 = self.up1(d2)
        d1 = self.dec1(torch.cat([d1, e1], dim=1))

        return self.out_conv(d1)

* 上記の実装は，元論文の U-Net 構成を単純化したものであり，論文の性能を再現するものではない．学習・検証を行った上で，層数やチャネル数の調整が必要になる場合がある．

学習と動作確認

ガウシアンノイズを付加した画像どうしのペアを用いて，モデルを学習させる．データセットは，任意の画像フォルダを使用する．

1. データセット用フォルダの準備

   学習用画像を格納するフォルダを作成する．

   cd /d c:%HOMEPATH%
   mkdir n2n_data
   cd n2n_data
   mkdir train
   

   train フォルダに，学習用の画像ファイル（jpg, png 等）を格納する．

2. 学習スクリプトの作成

   次のコードを train_n2n.py という名前で，n2n_model.py と同じフォルダに保存する．

   import argparse
   import glob
   import os
   
   import torch
   import torch.nn as nn
   from PIL import Image
   from torch.utils.data import Dataset, DataLoader
   from torchvision import transforms
   
   from n2n_model import N2NUNet
   
   
   class NoisyPairDataset(Dataset):
       """1枚の画像から，ノイズを加えた2枚のペアを作成するデータセット"""
   
       def __init__(self, image_dir, crop_size=128, noise_std=25):
           self.paths = sorted(
               glob.glob(os.path.join(image_dir, "*.jpg"))
               + glob.glob(os.path.join(image_dir, "*.png"))
           )
           self.crop_size = crop_size
           self.noise_std = noise_std
           self.transform = transforms.Compose(
               [
                   transforms.RandomCrop(
                       crop_size, pad_if_needed=True, padding_mode="reflect"
                   ),
                   transforms.ToTensor(),
               ]
           )
   
       def __len__(self):
           return len(self.paths)
   
       def __getitem__(self, idx):
           img = Image.open(self.paths[idx]).convert("RGB")
           clean = self.transform(img)
   
           noise1 = torch.randn_like(clean) * (self.noise_std / 255.0)
           noise2 = torch.randn_like(clean) * (self.noise_std / 255.0)
   
           noisy_input = (clean + noise1).clamp(0.0, 1.0)
           noisy_target = (clean + noise2).clamp(0.0, 1.0)
   
           return noisy_input, noisy_target
   
   
   def main():
       parser = argparse.ArgumentParser()
       parser.add_argument("--image_dir", type=str, required=True)
       parser.add_argument("--epochs", type=int, default=10)
       parser.add_argument("--batch_size", type=int, default=8)
       parser.add_argument("--crop_size", type=int, default=128)
       parser.add_argument("--lr", type=float, default=1e-3)
       parser.add_argument("--output_path", type=str, default="n2n_gaussian.pt")
       args = parser.parse_args()
   
       device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
       print(f"使用デバイス: {device}")
   
       dataset = NoisyPairDataset(args.image_dir, crop_size=args.crop_size)
       loader = DataLoader(
           dataset, batch_size=args.batch_size, shuffle=True, drop_last=True
       )
   
       model = N2NUNet().to(device)
       optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
       criterion = nn.MSELoss()
   
       for epoch in range(args.epochs):
           model.train()
           running_loss = 0.0
           for noisy_input, noisy_target in loader:
               noisy_input = noisy_input.to(device)
               noisy_target = noisy_target.to(device)
   
               optimizer.zero_grad()
               output = model(noisy_input)
               loss = criterion(output, noisy_target)
               loss.backward()
               optimizer.step()
   
               running_loss += loss.item()
   
           avg_loss = running_loss / len(loader)
           print(f"epoch {epoch + 1}/{args.epochs}  loss={avg_loss:.6f}")
   
       torch.save(model.state_dict(), args.output_path)
       print(f"モデルを保存した: {args.output_path}")
   
   
   if __name__ == "__main__":
       main()
   

3. 学習の実行

   コマンドプロンプトで，次のコマンドを実行する．

   cd /d c:%HOMEPATH%\n2n_data
   python train_n2n.py --image_dir train --epochs 10 --batch_size 8 --output_path n2n_gaussian.pt
   

4. テスト（ノイズ除去の実行）スクリプトの作成

   次のコードを test_n2n.py という名前で保存する．

   import argparse
   
   import torch
   from PIL import Image
   from torchvision import transforms
   
   from n2n_model import N2NUNet
   
   
   def main():
       parser = argparse.ArgumentParser()
       parser.add_argument("--weight_file", type=str, required=True)
       parser.add_argument("--input_image", type=str, required=True)
       parser.add_argument("--output_image", type=str, default="denoised.png")
       args = parser.parse_args()
   
       device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   
       model = N2NUNet().to(device)
       model.load_state_dict(torch.load(args.weight_file, map_location=device))
       model.eval()
   
       img = Image.open(args.input_image).convert("RGB")
       to_tensor = transforms.ToTensor()
       to_pil = transforms.ToPILImage()
   
       x = to_tensor(img).unsqueeze(0).to(device)
   
       with torch.no_grad():
           y = model(x).clamp(0.0, 1.0)
   
       result = to_pil(y.squeeze(0).cpu())
       result.save(args.output_image)
       print(f"ノイズ除去後の画像を保存した: {args.output_image}")
   
   
   if __name__ == "__main__":
       main()
   

5. テストの実行

   コマンドプロンプトで，次のコマンドを実行する．noisy_sample.jpg は，ノイズ除去を試したい画像ファイル名に置き換える．

   cd /d c:%HOMEPATH%\n2n_data
   python ..\test_n2n.py --weight_file n2n_gaussian.pt --input_image noisy_sample.jpg --output_image denoised.png