DCGANを用いた画像生成を自作データセットで行いたい

前提・実現したいこと

COlabでDCGANを用いて画像生成を行いたいと考えています。(プログラミング・AI初心者)
以下の記事を参考に，データセットの部分のみをオリジナルの画像(jpeg, 64×64 100枚)に置換したいのですが，データセットの作成で躓いています。
https://blog.shikoan.com/pytorch-dcgan/
作成するデータセットに用いたいファイルはgoogledriveにてjpegで保存しており，マウントは出来ています。

オリジナルデータセットの作り方のヒントだけでも教えていただけたら幸いです。

発生している問題・エラーメッセージ

NameError                                 Traceback (most recent call last)
<ipython-input-2-7e0c3ce0d650> in <module>()
    163 
    164 if __name__ == "__main__":
--> 165     train()

<ipython-input-2-7e0c3ce0d650> in train()
    101     for i in range(100):
    102         log_loss_G, log_loss_D = [], []
--> 103         for real_img, _ in tqdm(load_datasets()):
    104             batch_len = len(real_img)
    105 

NameError: name 'load_datasets' is not defined

該当のソースコード

pytorch
1
2import torch
3from torch import nn
4import torchvision
5import torchvision.transforms as transforms
6import numpy as np
7from tqdm import tqdm
8import os
9import pickle
10import statistics
11
12class Creat_datasets(Dataset):
13    def __init__(self, dir_name, csv_file, data_transform):
14        self.dir_name = dir_name
15        self.file = os.listdir("データセットにしたいフォルダのパス")
16        self.data_transform = data_transform
17         
18    def __len__(self):
19        return len(self.file)
20     
21    def __getitem__(self, i):
22         
23        image = Image.open(self.dir_name + self.file[i])
24        image = self.data_transform(image)
25        return image
26
27
28class Generator(nn.Module):
29    def __init__(self):
30        super().__init__()
31        self.main = nn.Sequential(
32            nn.ConvTranspose2d(64, 64, 4, 1, 0), # 4x4
33            nn.BatchNorm2d(64),
34            nn.ReLU(inplace=True),
35
36            nn.ConvTranspose2d(64, 32, 4, 1, 0), #7x7
37            nn.BatchNorm2d(32),
38            nn.ReLU(inplace=True),
39
40            nn.ConvTranspose2d(32, 16, 2, 2, 0), #14x14
41            nn.BatchNorm2d(16),
42            nn.ReLU(inplace=True),
43
44            nn.ConvTranspose2d(16, 1, 2, 2, 0), #28x28
45            nn.Tanh()
46        )
47
48    def forward(self, x):
49        return self.main(x)
50
51class Discriminator(nn.Module):
52    def __init__(self):
53        super().__init__()
54        self.model = nn.Sequential(
55            nn.Conv2d(1, 16, kernel_size=3, padding=1),
56            nn.BatchNorm2d(16),
57            nn.LeakyReLU(0.2, inplace=True),
58
59            nn.AvgPool2d(2),
60            nn.Conv2d(16, 32, kernel_size=3, padding=1),
61            nn.BatchNorm2d(32),
62            nn.LeakyReLU(0.2, inplace=True),
63
64            nn.AvgPool2d(2),
65            nn.Conv2d(32, 64, kernel_size=3, padding=1),
66            nn.BatchNorm2d(64),
67            nn.LeakyReLU(0.2, inplace=True),
68
69            nn.AvgPool2d(7),
70            nn.Conv2d(64, 1, 1) # fcの代わり
71        )
72
73    def forward(self, x):
74        return self.model(x).squeeze()
75
76
77def train():
78    # モデル
79    device = "cuda"
80    model_G, model_D = Generator(), Discriminator()
81    model_G, model_D = nn.DataParallel(model_G), nn.DataParallel(model_D)
82    model_G, model_D = model_G.to(device), model_D.to(device)
83
84    params_G = torch.optim.Adam(model_G.parameters(),
85                lr=0.0002, betas=(0.5, 0.999))
86    params_D = torch.optim.Adam(model_D.parameters(),
87                lr=0.0002, betas=(0.5, 0.999))
88    
89    # ロスを計算するためのラベル変数
90    ones = torch.ones(512).to(device)
91    zeros = torch.zeros(512).to(device)
92    loss_f = nn.BCEWithLogitsLoss()
93
94    # エラー推移
95    result = {}
96    result["log_loss_G"] = []
97    result["log_loss_D"] = []
98
99    # 訓練
100    for i in range(100):
101        log_loss_G, log_loss_D = [], []
102        for real_img, _ in tqdm(load_datasets()):
103            batch_len = len(real_img)
104
105            # Gの訓練
106            # 偽画像を作成
107            z = torch.randn(batch_len, 64, 1, 1).to(device)
108            fake_img = model_G(z)
109
110            # 偽画像を一時保存
111            fake_img_tensor = fake_img.detach()
112
113            # 偽画像を本物と騙せるようにロスを計算
114            out = model_D(fake_img)
115            loss_G = loss_f(out, ones[:batch_len])
116            log_loss_G.append(loss_G.item())
117
118            # 微分計算・重み更新
119            params_D.zero_grad()
120            params_G.zero_grad()
121            loss_G.backward()
122            params_G.step()
123
124            # Discriminatoの訓練
125            # sample_dataの実画像
126            real_img = real_img.to(device)
127            # 実画像を実画像と識別できるようにロスを計算
128            real_out = model_D(real_img)
129            loss_D_real = loss_f(real_out, ones[:batch_len])
130
131            # 偽の画像の偽と識別できるようにロスを計算
132            fake_out = model_D(fake_img_tensor)
133            loss_D_fake = loss_f(fake_out, zeros[:batch_len])
134
135            # 実画像と偽画像のロスを合計
136            loss_D = loss_D_real + loss_D_fake
137            log_loss_D.append(loss_D.item())
138
139            # 微分計算・重み更新
140            params_D.zero_grad()
141            params_G.zero_grad()
142            loss_D.backward()
143            params_D.step()
144
145        result["log_loss_G"].append(statistics.mean(log_loss_G))
146        result["log_loss_D"].append(statistics.mean(log_loss_D))
147        print("log_loss_G =", result["log_loss_G"][-1], ", log_loss_D =", result["log_loss_D"][-1])
148
149        # 画像を保存
150        if not os.path.exists("mnist_generated"):
151            os.mkdir("mnist_generated")
152        torchvision.utils.save_image(fake_img_tensor[:min(batch_len, 100)],
153                                f"mnist_generated/epoch_{i:03}.png")
154    # ログの保存
155    with open("mnist_generated/logs.pkl", "wb") as fp:
156        pickle.dump(result, fp)
157
158def save_test():
159    for X, _ in load_datasets():
160        torchvision.utils.save_image(X[:100], "mnist_generated/true.png")
161        break
162
163if __name__ == "__main__":
164    train()

試したこと

MNISTの定義部分をまるごと取り替えてしまえばよいのかと思い，元記事のMNISTのデータセット作成のコードをこちらの記事(https://betashort-lab.com/%E3%83%87%E3%83%BC%E3%82%BF%E3%82%B5%E3%82%A4%E3%82%A8%E3%83%B3%E3%82%B9/%E3%83%87%E3%82%A3%E3%83%BC%E3%83%97%E3%83%A9%E3%83%BC%E3%83%8B%E3%83%B3%E3%82%B0/pytorch%E3%81%AEdatasets%E3%81%A7%E7%94%BB%E5%83%8F%E3%83%87%E3%83%BC%E3%82%BF%E3%82%BB%E3%83%83%E3%83%88%E3%82%92%E4%BD%9C%E3%82%8B/)を参考にして，

class Creat_datasets(Dataset):
def init(self, dir_name, csv_file, data_transform):
self.dir_name = dir_name
self.file = os.listdir("データセットにしたいフォルダのパス")
self.data_transform = data_transform

def __len__(self):
    return len(self.file)
 
def __getitem__(self, i):
     
    image = Image.open(self.dir_name + self.file[i])
    image = self.data_transform(image)
    return image

に書き換えました。
その結果，上記のエラーが発生しています。

補足情報（FW/ツールのバージョンなど）

ここにより詳細な情報を記載してください。

jbpb0

2023/02/01 00:29 編集

https://blog.shikoan.com/pytorch-dcgan/ の「コード」の「def load_datasets():」を削除したら、 > NameError: name 'load_datasets' is not defined となるのは当然です「def load_datasets():」を削除しないで、それの「trainset = torchvision.datasets.CIFAR10(...」の行を、 trainset = torchvision.datasets.ImageFolder(root="データセットにしたいフォルダのパス", transform=transform) みたいにしてみたら、いかがでしょうか参考 https://qiita.com/coticoticotty/items/5fe4a7f3c740523eb2fb https://reafnex.net/ai/pytorch-use-imagefolder/ の「ImageFolderをつかってみる」