THINGS datasetというのを使って人間の判断に類似するように画像認識モデルをファインチューンできるらしい。
https://arxiv.org/abs/2409.06509https://doi.org/10.1038/s41586-025-09631-6 (nature)
~ただのメモ置き場。なので、内容の正しさについての保証は一切ありませんよ~ copyright (c) 2006-2025 bluewidz, all rights reserved.
THINGS datasetというのを使って人間の判断に類似するように画像認識モデルをファインチューンできるらしい。
https://arxiv.org/abs/2409.06509
投稿者
bluewidz
0
コメント
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | import matplotlib.pyplot as plt import numpy as np import random import torch import torch.nn as nn import torch.nn.functional as F from torchvision import datasets, transforms from torch.utils.data import DataLoader from pykan.kan.KANLayer import KANLayer import sys def initialize_seed(seed=0): torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) class ConvMLP(nn.Module): def __init__(self, fc_layers: list[int], device): super().__init__() self.conv1 = nn.Conv2d(1, 8, kernel_size=5, stride=2, device=device) self.conv2 = nn.Conv2d(8, 16, kernel_size=5, stride=2, device=device) n_in = 16 fcs = [] for fc in fc_layers: fcs.append(nn.Linear(n_in, fc, device=device)) n_in = fc self.fcs = nn.ModuleList(fcs) def forward(self, x): x = self.conv2(F.relu(F.max_pool2d(self.conv1(x), 2))) x = x.reshape(x.shape[0], -1) for fc in self.fcs: x = fc(F.relu(x)) return x def update_grid_from_samples(self, x): pass def regularize(self, lambda_l1, lambda_entropy, lambda_coef, lambda_coefdiff, small_mag_threshold=1e-16, small_reg_factor=1.0): return 0.0 # Modified version of KAN in pykan/kan/KAN.py class ConvKAN(nn.Module): def __init__(self, width: list[int], grid=5, k=3, noise_scale=0.1, noise_scale_base=0.1, base_fun=torch.nn.SiLU(), bias_trainable=True, grid_eps=1.0, grid_range=[-1, 1], sp_trainable=True, sb_trainable=True, device="cpu"): super().__init__() ### Initialize feature extraction layers self.conv1 = nn.Conv2d(1, 8, kernel_size=5, stride=2, device=device) self.conv2 = nn.Conv2d(8, 16, kernel_size=5, stride=2, device=device) width.insert(0, 16) ### Initialize KAN layers self.biases = [] self.act_fun = [] self.depth = len(width) - 1 self.width = width for l in range(self.depth): # splines scale_base = 1 / np.sqrt(width[l]) + (torch.randn(width[l] * width[l + 1], ) * 2 - 1) * noise_scale_base sp_batch = KANLayer(in_dim=width[l], out_dim=width[l + 1], num=grid, k=k, noise_scale=noise_scale, scale_base=scale_base, scale_sp=1.0, base_fun=base_fun, grid_eps=grid_eps, grid_range=grid_range, sp_trainable=sp_trainable, sb_trainable=sb_trainable, device=device) self.act_fun.append(sp_batch) # bias bias = nn.Linear(width[l + 1], 1, bias=False, device=device).requires_grad_(bias_trainable) bias.weight.data *= 0.0 self.biases.append(bias) self.biases = nn.ModuleList(self.biases) self.act_fun = nn.ModuleList(self.act_fun) def forward(self, x): # Extract features by conv x = self.conv2(F.relu(F.max_pool2d(self.conv1(x), 2))) x = x.reshape(x.shape[0], -1) # Run KAN layers self.acts = [x] # acts shape: (batch, width[l]) self.acts_scale = [] for l in range(self.depth): x, preacts, postacts, postspline = self.act_fun[l](x) grid_reshape = self.act_fun[l].grid.reshape(self.width[l + 1], self.width[l], -1) input_range = grid_reshape[:, :, -1] - grid_reshape[:, :, 0] + 1e-4 output_range = torch.mean(torch.abs(postacts), dim=0) self.acts_scale.append(output_range / input_range) x = x + self.biases[l].weight self.acts.append(x) return x def update_grid_from_samples(self, x): for l in range(self.depth): self.forward(x) self.act_fun[l].update_grid_from_samples(self.acts[l]) def regularize(self, lambda_l1, lambda_entropy, lambda_coef, lambda_coefdiff, small_mag_threshold=1e-16, small_reg_factor=1.0): def nonlinear(x, th, factor): return (x < th) * x * factor + (x > th) * (x + (factor - 1) * th) reg_ = 0. for i in range(len(self.acts_scale)): vec = self.acts_scale[i].reshape(-1, ) vec_sum = torch.sum(vec) if vec_sum == 0.0: continue p = vec / vec_sum l1 = torch.sum(nonlinear(vec, th=small_mag_threshold, factor=small_reg_factor)) entropy = - torch.sum(p * torch.log2(p + 1e-4)) reg_ += lambda_l1 * l1 + lambda_entropy * entropy # both l1 and entropy # regularize coefficient to encourage spline to be zero for i in range(len(self.act_fun)): coeff_l1 = torch.sum(torch.mean(torch.abs(self.act_fun[i].coef), dim=1)) coeff_diff_l1 = torch.sum(torch.mean(torch.abs(torch.diff(self.act_fun[i].coef)), dim=1)) reg_ += lambda_coef * coeff_l1 + lambda_coefdiff * coeff_diff_l1 return reg_ def calc_accuracy(ys): rs = [] for y, label in ys: rs.append((torch.argmax(y, dim=1) == label).float()) r = torch.cat(rs, dim=0) return torch.mean(r)*100.0 def train(model, train_loader, test_loader, max_epoch, lamb=0.0, lambda_l1=1.0, lambda_entropy=2.0, lambda_coef=0.0, lambda_coefdiff=0.0, update_grid=True, grid_update_freq=10, loss_fn=torch.nn.CrossEntropyLoss(), lr=0.002, device="cpu"): optimizer = torch.optim.Adam(model.parameters(), lr=lr) for epoch in range(max_epoch): model.train() n_samples = 0 max_samples = len(train_loader.dataset) ys = [] for iter, (x, label) in enumerate(train_loader): x = x.to(device) label = label.to(device) if iter % grid_update_freq == 0 and update_grid: model.update_grid_from_samples(x) y = model(x) ys.append((y, label)) loss = loss_fn(y, label) reg_ = model.regularize(lambda_l1, lambda_entropy, lambda_coef, lambda_coefdiff) loss = loss + lamb * reg_ optimizer.zero_grad() loss.backward() optimizer.step() n_samples += len(x) if iter % 100 == 0: print(f"Epoch: {epoch} [{n_samples}/{max_samples}] Loss: {loss.item():.6f}") # Calc train accuracy train_acc = calc_accuracy(ys) # Calc test accuracy model.eval() ys = [] with torch.no_grad(): for iter, (x, label) in enumerate(test_loader): x = x.to(device) label = label.to(device) y = model(x) ys.append((y, label)) test_acc = calc_accuracy(ys) print(f"Epoch: {epoch} [{n_samples}/{max_samples}] Loss: {loss.item():.6f} Acc(Train): {train_acc} Acc(Test): {test_acc}") return def main(mode): initialize_seed(123) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") reg_lambda = 0.0 update_grid = True, if mode == "kan": model = ConvKAN(width=[20, 10], device=device) elif mode == "kan-no-update-grid": model = ConvKAN(width=[20, 10], device=device) update_grid = False elif mode == "kan-reg": model = ConvKAN(width=[20, 10], device=device) reg_lambda = 0.003 elif mode == "mlp": model = ConvMLP(fc_layers=[20, 10], device=device) else: return train_loader = DataLoader(datasets.MNIST("./data", train=True, download=True, transform=transforms.ToTensor()), batch_size=128, shuffle=True) test_loader = DataLoader(datasets.MNIST("./data", train=False, download=True, transform=transforms.ToTensor()), batch_size=128, shuffle=False) train(model, train_loader, test_loader, max_epoch=5, lamb=reg_lambda, update_grid=update_grid, device=device) if __name__ == "__main__": main(sys.argv[1]) |
Epoch: 0 [128/60000] Loss: 2.330844 Epoch: 0 [12928/60000] Loss: 0.527279 Epoch: 0 [25728/60000] Loss: 0.264145 Epoch: 0 [38528/60000] Loss: 0.268547 Epoch: 0 [51328/60000] Loss: 0.270200 Epoch: 0 [60000/60000] Loss: 0.137060 Acc(Train): 82.13166809082031 Acc(Test): 93.66999816894531 Epoch: 1 [128/60000] Loss: 0.211759 Epoch: 1 [12928/60000] Loss: 0.172556 Epoch: 1 [25728/60000] Loss: 0.216685 Epoch: 1 [38528/60000] Loss: 0.224366 Epoch: 1 [51328/60000] Loss: 0.166860 Epoch: 1 [60000/60000] Loss: 0.297034 Acc(Train): 94.53333282470703 Acc(Test): 95.87000274658203 Epoch: 2 [128/60000] Loss: 0.169406 Epoch: 2 [12928/60000] Loss: 0.060187 Epoch: 2 [25728/60000] Loss: 0.065257 Epoch: 2 [38528/60000] Loss: 0.209189 Epoch: 2 [51328/60000] Loss: 0.137371 Epoch: 2 [60000/60000] Loss: 0.062734 Acc(Train): 95.69499969482422 Acc(Test): 96.29000091552734 Epoch: 3 [128/60000] Loss: 0.161130 Epoch: 3 [12928/60000] Loss: 0.094211 Epoch: 3 [25728/60000] Loss: 0.137475 Epoch: 3 [38528/60000] Loss: 0.143321 Epoch: 3 [51328/60000] Loss: 0.086744 Epoch: 3 [60000/60000] Loss: 0.295325 Acc(Train): 96.2300033569336 Acc(Test): 96.52999877929688 Epoch: 4 [128/60000] Loss: 0.084705 Epoch: 4 [12928/60000] Loss: 0.114923 Epoch: 4 [25728/60000] Loss: 0.071916 Epoch: 4 [38528/60000] Loss: 0.093224 Epoch: 4 [51328/60000] Loss: 0.111265 Epoch: 4 [60000/60000] Loss: 0.049419 Acc(Train): 96.69833374023438 Acc(Test): 97.1500015258789
Epoch: 0 [128/60000] Loss: 2.308975 Epoch: 0 [12928/60000] Loss: 0.608790 Epoch: 0 [25728/60000] Loss: 0.333690 Epoch: 0 [38528/60000] Loss: 0.214927 Epoch: 0 [51328/60000] Loss: 0.171603 Epoch: 0 [60000/60000] Loss: 0.099660 Acc(Train): 88.74166870117188 Acc(Test): 96.1500015258789 Epoch: 1 [128/60000] Loss: 0.073898 Epoch: 1 [12928/60000] Loss: 0.215137 Epoch: 1 [25728/60000] Loss: 0.109934 Epoch: 1 [38528/60000] Loss: 0.126619 Epoch: 1 [51328/60000] Loss: 0.066091 Epoch: 1 [60000/60000] Loss: 0.188135 Acc(Train): 96.20833587646484 Acc(Test): 96.88999938964844 Epoch: 2 [128/60000] Loss: 0.061085 Epoch: 2 [12928/60000] Loss: 0.078620 Epoch: 2 [25728/60000] Loss: 0.045636 Epoch: 2 [38528/60000] Loss: 0.052172 Epoch: 2 [51328/60000] Loss: 0.037537 Epoch: 2 [60000/60000] Loss: 0.149782 Acc(Train): 96.87333679199219 Acc(Test): 95.19000244140625 Epoch: 3 [128/60000] Loss: 0.040984 Epoch: 3 [12928/60000] Loss: 0.102282 Epoch: 3 [25728/60000] Loss: 0.017132 Epoch: 3 [38528/60000] Loss: 0.043684 Epoch: 3 [51328/60000] Loss: 0.126490 Epoch: 3 [60000/60000] Loss: 0.057794 Acc(Train): 97.1483383178711 Acc(Test): 97.47000122070312 Epoch: 4 [128/60000] Loss: 0.056742 Epoch: 4 [12928/60000] Loss: 0.087390 Epoch: 4 [25728/60000] Loss: 0.046712 Epoch: 4 [38528/60000] Loss: 0.058726 Epoch: 4 [51328/60000] Loss: 0.217805 Epoch: 4 [60000/60000] Loss: 0.137546 Acc(Train): 97.54500579833984 Acc(Test): 97.52999877929688
Epoch: 0 [128/60000] Loss: 2.308975 Epoch: 0 [12928/60000] Loss: 0.502888 Epoch: 0 [25728/60000] Loss: 0.330806 Epoch: 0 [38528/60000] Loss: 0.242808 Epoch: 0 [51328/60000] Loss: 0.165698 Epoch: 0 [60000/60000] Loss: 0.191306 Acc(Train): 86.51166534423828 Acc(Test): 95.06999969482422 Epoch: 1 [128/60000] Loss: 0.090638 Epoch: 1 [12928/60000] Loss: 0.238503 Epoch: 1 [25728/60000] Loss: 0.128466 Epoch: 1 [38528/60000] Loss: 0.166372 Epoch: 1 [51328/60000] Loss: 0.120421 Epoch: 1 [60000/60000] Loss: 0.113729 Acc(Train): 95.87166595458984 Acc(Test): 96.37999725341797 Epoch: 2 [128/60000] Loss: 0.094055 Epoch: 2 [12928/60000] Loss: 0.101774 Epoch: 2 [25728/60000] Loss: 0.053376 Epoch: 2 [38528/60000] Loss: 0.050028 Epoch: 2 [51328/60000] Loss: 0.049892 Epoch: 2 [60000/60000] Loss: 0.045015 Acc(Train): 96.88333129882812 Acc(Test): 96.44000244140625 Epoch: 3 [128/60000] Loss: 0.091558 Epoch: 3 [12928/60000] Loss: 0.116339 Epoch: 3 [25728/60000] Loss: 0.035658 Epoch: 3 [38528/60000] Loss: 0.047689 Epoch: 3 [51328/60000] Loss: 0.121902 Epoch: 3 [60000/60000] Loss: 0.066078 Acc(Train): 97.5 Acc(Test): 97.30999755859375 Epoch: 4 [128/60000] Loss: 0.069760 Epoch: 4 [12928/60000] Loss: 0.048372 Epoch: 4 [25728/60000] Loss: 0.042325 Epoch: 4 [38528/60000] Loss: 0.073130 Epoch: 4 [51328/60000] Loss: 0.130568 Epoch: 4 [60000/60000] Loss: 0.046441 Acc(Train): 97.72833251953125 Acc(Test): 97.43999481201172
Epoch: 0 [128/60000] Loss: 2.449388 Epoch: 0 [12928/60000] Loss: 0.862422 Epoch: 0 [25728/60000] Loss: 0.502983 Epoch: 0 [38528/60000] Loss: 0.569132 Epoch: 0 [51328/60000] Loss: 0.391648 Epoch: 0 [60000/60000] Loss: 0.360877 Acc(Train): 89.41166687011719 Acc(Test): 95.72999572753906 Epoch: 1 [128/60000] Loss: 0.345524 Epoch: 1 [12928/60000] Loss: 0.449618 Epoch: 1 [25728/60000] Loss: 0.375242 Epoch: 1 [38528/60000] Loss: 0.334639 Epoch: 1 [51328/60000] Loss: 0.313511 Epoch: 1 [60000/60000] Loss: 0.374076 Acc(Train): 96.25166320800781 Acc(Test): 96.95999908447266 Epoch: 2 [128/60000] Loss: 0.237053 Epoch: 2 [12928/60000] Loss: 0.721758 Epoch: 2 [25728/60000] Loss: 0.468227 Epoch: 2 [38528/60000] Loss: 0.489218 Epoch: 2 [51328/60000] Loss: 0.366972 Epoch: 2 [60000/60000] Loss: 0.461017 Acc(Train): 90.59166717529297 Acc(Test): 94.61000061035156 Epoch: 3 [128/60000] Loss: 0.391373 Epoch: 3 [12928/60000] Loss: 0.405158 Epoch: 3 [25728/60000] Loss: 0.296408 Epoch: 3 [38528/60000] Loss: 0.306373 Epoch: 3 [51328/60000] Loss: 0.362922 Epoch: 3 [60000/60000] Loss: 0.335211 Acc(Train): 94.8933334350586 Acc(Test): 95.55999755859375 Epoch: 4 [128/60000] Loss: 0.343693 Epoch: 4 [12928/60000] Loss: 0.300828 Epoch: 4 [25728/60000] Loss: 0.310633 Epoch: 4 [38528/60000] Loss: 0.402329 Epoch: 4 [51328/60000] Loss: 0.401533 Epoch: 4 [60000/60000] Loss: 0.272522 Acc(Train): 95.64000701904297 Acc(Test): 95.79000091552734\(\lambda=0.005\)にすると、以下のようになり途中でモデルが崩壊しました。
Epoch: 0 [128/60000] Loss: 2.542997 Epoch: 0 [12928/60000] Loss: 0.992837 Epoch: 0 [25728/60000] Loss: 0.666648 Epoch: 0 [38528/60000] Loss: 0.642962 Epoch: 0 [51328/60000] Loss: 0.496544 Epoch: 0 [60000/60000] Loss: 0.467141 Acc(Train): 88.85499572753906 Acc(Test): 95.80999755859375 Epoch: 1 [128/60000] Loss: 0.465318 Epoch: 1 [12928/60000] Loss: 0.551555 Epoch: 1 [25728/60000] Loss: 0.480575 Epoch: 1 [38528/60000] Loss: 0.532717 Epoch: 1 [51328/60000] Loss: 0.389326 Epoch: 1 [60000/60000] Loss: 0.464731 Acc(Train): 95.5183334350586 Acc(Test): 96.31999969482422 Epoch: 2 [128/60000] Loss: 0.360698 Epoch: 2 [12928/60000] Loss: 0.528383 Epoch: 2 [25728/60000] Loss: 0.377956 Epoch: 2 [38528/60000] Loss: 0.333951 Epoch: 2 [51328/60000] Loss: 0.379329 Epoch: 2 [60000/60000] Loss: 0.398343 Acc(Train): 94.8566665649414 Acc(Test): 96.06999969482422 Epoch: 3 [128/60000] Loss: 0.374904 Epoch: 3 [12928/60000] Loss: 0.928802 Epoch: 3 [25728/60000] Loss: 0.401091 Epoch: 3 [38528/60000] Loss: 0.539629 Epoch: 3 [51328/60000] Loss: 223.694611 Epoch: 3 [60000/60000] Loss: 3.871729 Acc(Train): 76.33499908447266 Acc(Test): 10.520000457763672 Epoch: 4 [128/60000] Loss: 4.122091 Epoch: 4 [12928/60000] Loss: 3.226183 Epoch: 4 [25728/60000] Loss: 2.973773 Epoch: 4 [38528/60000] Loss: 2.831389 Epoch: 4 [51328/60000] Loss: 2.789084 Epoch: 4 [60000/60000] Loss: 2.598794 Acc(Train): 9.819999694824219 Acc(Test): 10.09999942779541
投稿者
bluewidz
0
コメント
class ConditionalUNet(UNet):
(省略)
def forward(self, x: Tensor, t: Tensor, label: Tensor) -> Tensor:
x0 = x #self.embedding1(x, t)
x1 = self.block1(x0)
x1 = self.label_emb1(x1, label)
#x1 = self.embedding2(x1, t)
x2 = self.block2(self.down1(x1))
x2 = self.label_emb2(x2, label)
#x2 = self.embedding3(x2, t)
crossed = self.label_emb3(self.block3(self.down2(x2)), label)
x3 = self.up1(self.attention1(crossed))
x4 = torch.cat([x2, x3], dim=1)
#x4 = self.embedding4(x4, t)
x5 = self.up2(self.label_emb4(self.block4(x4), label))
x6 = torch.cat([x5, x1], dim=1)
x6 = self.label_emb5(x6, label)
#x6 = self.embedding5(x6, t)
out = self.out(self.block5(x6))
return out
とします。
投稿者
bluewidz
0
コメント
variance_schedule = torch.linspace(1e-4, 0.01, steps=1000)となっています。これをパターンAでは
variance_schedule = torch.linspace(1e-4, 0.1, steps=100)と、パターンBでは
variance_schedule = torch.linspace(1e-4, 0.999, steps=10)とします。 それぞれのスケジュールを使ったときのalphaは
[パターン デフォルト] [0.99990, 0.99989, 0.99988, ... , 0.99001, 0.99000] [パターン A] [0.99990, 0.99889, 0.99788, ... , 0.90101, 0.90000] [パターン B] [0.99990, 0.88891, 0.77792, ... , 0.11199, 0.00100]となります。 alpha_barは
[パターン デフォルト] [0.9999, 0.9998, 0.9997, ... , 0.0064, 0.0063] [パターン A] [0.9999, 0.9988, 0.9967, ... , 0.0062, 0.0056] [パターン B] [9.9990e-01, 8.8882e-01, 6.9143e-01, ... , 9.5131e-04, 9.5130e-07]となります。ここで重要なことは、最初の時刻(ノイズが乗っていない)をt=0、最後の時刻(完全にノイズ)をt=1とするとき、alpha_barはt=0では1に近く、t=1では0に近くなるようにvariance_scheduleを決める必要があるということです。 各時刻tにおけるノイズの強さがalpha_barで決まり、t=1のときに完全にノイズになっていないと拡散プロセスの前提が崩れてしまうためです。 実際、パターンAの
variance_schedule = torch.linspace(1e-4, 0.1, steps=100)を
variance_schedule = torch.linspace(1e-4, 0.01, steps=100)に変えると、alpha_bar の値は
0.99990, 0.99970, 0.99940, ... , 0.60857, 0.60248となりますが、この場合、数字の画像をうまく生成できません。
t = torch.randint(0, 1000, (image.shape[0],))の1000のところを、パターンAでは100に、パターンBでは10にします。
for t in trange(999, -1, -1):の999のところを、パターンAでは99に、パターンBでは9にします。刻む数が少なくなると(ステップの細かさを粗くすると)、その分だけ生成時間を短くできます。
投稿者
bluewidz
0
コメント
device = "cpu"と書かれている行を
device = "cuda"に書き換えます。そして
$ python main.pyを実行すると、学習が始まります。しばらく待っていると完了し、モデルのパラメータがunet_mnist.pthに記録されます。
- device = "cpu" + device = "cuda"に書き換えます。-が書き換え対象の行で、+が書き換えたあとの行の内容です。さらに、
- from models import UNet + from models import UNet, ConditionalUNetと書き換え、
- model = UNet().to(device) + model = ConditionalUNet().to(device)と書き換えます。ConditionalUNetにすると、生成時にどの数字を生成するかを指定する必要があるため、
model.eval()
+ labels = torch.randint(0, 10, (batch_size,))
+ print(labels)
+ labels_cpu = labels
+ labels = labels.to(device)
with torch.no_grad():
for t in trange(999, -1, -1):
time = torch.ones(batch_size) * t
- et = model(xt.to(device), time.to(device)) # predict noise
+ et = model(xt.to(device), time.to(device), labels) # predict noise
xt = process.inverse(xt, et.cpu(), t)
のように、書き換えます。ランダムに0〜9の値をラベルとして指定するようにしています。
出力部分を少し書き換えて
labels = ["Generated Images"] * 9
-
- for i in range(9):
- plt.subplot(3, 3, i + 1)
- plt.tight_layout()
- plt.imshow(xt[i][0], cmap="gray", interpolation="none")
- plt.title(labels[i])
- plt.show()
+ if t % 10 == 0:
+ plt.figure(figsize=(10, 10))
+ for i in range(25):
+ plt.subplot(5, 5, i + 1)
+ plt.tight_layout()
+ plt.imshow(xt[i][0], cmap="gray", interpolation="none")
+ plt.title(f"{labels_cpu[i]}")
+ plt.savefig(f"images/generated_t{t}.png")
+ #plt.show()
とすると、逆拡散過程によって少しずつ数字が画像として浮かび上がるところを見ることができます。
ただし、このようにすると、生成した数字のpyplotでの描画処理のために生成時間が延びます。GPU使用率もあきらかに低下します。
途中結果を見る必要がなければ最初のコードのほうが良いでしょう。
for epoch in range(100):
for image, label in 学習用画像とそのラベルの集合:
# 画像に加えるノイズの強さをランダムに選択
t = torch.randint(0, 1000, (image.shape[0],))
# 画像に加えるノイズを作成
epsilon = torch.randn(image.shape)
# tとepsilonに基づいてノイズをimageに加える
diffused_image = process.forward(image, t, epsilon)
# modelを使って加えられたノイズを予測
optimizer.zero_grad()
output = model(diffused_image, t, label)
# 予測したノイズがどれだけ正しいかを評価して、モデルの重みを更新
loss = criterion(epsilon, output) # criterion = torch.nn.MSELoss()
loss.backward()
optimizer.step()
となります。ノイズが加えられた画像から、加えられたノイズを推定するモデルを学習しているだけです。
それ以外は通常のニューラルネットワークの学習と変わりがありません。
そしてこれは元論文のAlgorithm 1の通りの実装です。5行目の
と見比べると、\(\bm{\epsilon}\)はコード上のepsilonに、 \(\bm{\epsilon_{\theta}}\)はコード上のmodelに、 \(\sqrt{\bar{\alpha_t}}\bm{x}_0 + \sqrt{1-\bar{\alpha}_t}\bm{\epsilon}\)はdiffused_imageに相当していることが分かります。 Algorithm 1とは異なり、modelの引数にlabelが余計についていますが、これは生成する数字が0〜9のどれであるかをニューラルネットワークに指示するための値となります。
\( \nabla_\theta \| \bm{\epsilon} - \bm{\epsilon}_\theta ( \sqrt{\bar{\alpha_t}}\bm{x}_0 + \sqrt{1-\bar{\alpha}_t}\bm{\epsilon}, t) \|^2 \)
diffused_image = process.forward(image, t, epsilon)の部分について、詳細を見ていきます。 まず、diffusion_model.pyのDiffusionProcessの初期化処理部分である__init__にて\(\alpha_t\)の値が計算されています。 最初に\(\beta_t\) (コード上ではvariance_schedule) を
self.variance_schedule = torch.linspace(1e-4, 0.01, steps=1000)のように計算しています。具体的には
self.variance_schedule = [1.0000e-04, 1.0991e-04, 1.1982e-04, ... , 0.009980, 0.009990, 0.01]となっています。元論文では\(\beta_1=10^4\)、\(\beta_T=0.02\)、\(T=1000\)と書かれているので\(\beta_T\)の値のみ異なっています。 \(\alpha_t = 1-\beta_t\)であるので、
self.alpha = 1 - self.variance_scheduleにて計算され、具体的な値は
self.alpha = [0.999900, 0.999890, 0.999880 , ... , 0.990020, 0.990010, 0.990000]となります。さらに、\(\bar{\alpha}_t=\prod_{s=1}^t \alpha_s\)であるので、
self.alpha_bar = torch.cumprod(self.alpha, dim=0)にて計算され、具体的な値は
self.alpha_bar = [0.999900, 0.999790, 0.999670, ... , 0.006430, 0.006365, 0.006302]となります。 これらの値を使って、forwardにてノイズを画像に付加します。 \(\sqrt{\bar{\alpha_t}}\bm{x}_0 + \sqrt{1-\bar{\alpha}_t}\bm{\epsilon}\)が計算できれば良いので、まず\(\sqrt{1-\bar{\alpha}_t}\)を
std_dev = torch.sqrt(1 - self.alpha_bar[time_step])で計算します。次に\(\sqrt{\bar{\alpha_t}}\)を
mean_multiplier = torch.sqrt(self.alpha_bar[time_step])で計算し、最後に
diffused_images = mean_multiplier * x_0 + std_dev * noiseのように足し合わせます。x_0が\(\bm{x}_0\)で、noiseが\(\bm{\epsilon}\)です。
half_dim = self.dim // 2で\(d/2\)を計算し、
embeddings = math.log(10000) / (half_dim - 1)で、\(\frac{1}{\frac{d}{2}-1} \log 10000\)を計算し、続いて
embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)で、\(\exp \left(-\frac{i}{\frac{d}{2}-1} \log 10000 \right)\)を計算していることが分かります。 なお、変形前の\(1/10000^{i/(d/2-1)}\)にしたがって
1/torch.pow(10000, torch.arange(half_dim, device=device)/(half_dim-1))で計算しても、結果はほとんど変わりません。具体的には、最初の5つの値を出力すると、
式変形前 [1.0000000000, 0.9821373820, 0.9645937681, 0.9473634958, 0.9304410219] 式変形後 [1.0000000000, 0.9821373224, 0.9645937085, 0.9473634958, 0.9304410219]となるので、どちらで計算してもよさそうです。 さて、時刻\(t\)をかけると
embeddings = time[:, None] * embeddings[None, :]となり、最後に、
embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)のように\(sin\)と\(cos\)を通し、それらを連結することで、埋め込みベクトルが作成できます。
for t in trange(999, -1, -1):
# 時刻tとラベルlabelの条件の下、ノイズを推定する
et = model(xt, t, label)
# 推定したノイズを使って画像からノイズを除去する
xt = process.inverse(xt, et, t)
# xtを画像として描画する
draw(xt)
となります。モデルの学習時は時刻tをランダムに選んでいましたが、
生成するときはノイズから徐々にノイズを取り除いていきます。t=999ではノイズのみ、t=0ではノイズがなくなった画像になります。
xt = process.inverse(xt, et, t)の部分について詳細を見ていきます。 元論文arXiv:2006.11239のAlgorithm 2の4行目の処理 \[ \bm{x}_{t-1}=\frac{1}{\sqrt{\alpha_t}}\left(\bm{x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(\bm{x}_t,t)\right)+\sigma_t\bm{z} \] が実装されています。DiffusionProcessのinverseを見ていくと、
scale = 1 / torch.sqrt(self.alpha[t])では\(\frac{1}{\sqrt{\alpha_t}}\)の部分が計算されています。
noise_scale = (1 - self.alpha[t]) / torch.sqrt(1 - self.alpha_bar[t])では\(\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\)の部分が計算されています。
std_dev = torch.sqrt(self.variance_schedule[t])では\(\sigma_t\)が計算されています。
z = torch.randn(xt.shape) if t > 1 else torch.Tensor([0])では\(\bm{z} \sim N(0,\bm{I})\)が計算されています。ガウス分布に従うノイズを作っているだけです。 最後にすべてをつなげて
mu_t = scale * (xt - noise_scale * et) xt = mu_t + std_dev * z # remove noise from imageを計算することで、Algorithm 2の4行目の処理を実現しています。
投稿者
bluewidz
0
コメント
いまさらではあるものの、Self-attention (自己注意)の計算方法について途中の行列の形状に着目して調べてみました。
結果だけ知りたい方は、最後のまとめに進んでください。
def compute_qkv(query_antecedent, ←これがcommon_kayers.layer_preprocess(x, hparams)
memory_antecedent, ←これがNone
total_key_depth,
total_value_depth,
...
となっていて、memory_antecedentがNoneであることは、呼び出し元のtransformer_encoder:213に戻るとわかります。
memory_antecedentがNoneならquery_antecedentにしているので、self-attentionの場合、入力はquery_antecedentと考えればよさそうです。
さて、Q, K, V (コード中ではq, k, v)はcompute_attention_componentで計算されていて、入力は
antecedent: a Tensor with shape [batch, length, channels]戻り値が
c : [batch, length, depth] tensorとなっています。filter_widthによって処理内容が異なるようですが、filter_width == 1のケースを見てみると、 4415行目で
return common_layers.dense(
antecedent, total_depth, use_bias=False, name=name,
layer_collection=layer_collection)
のように書かれています。バイアスなしなので、単にMatMulの計算をしているだけになります。
Tensorflowのドキュメントによると、
Dense implements the operation: output = activation(dot(input, kernel) + bias) <中略> kernel is a weights matrix created by the layer,なので、
q = (batch, length_q, total_key_depth) k = (batch, length_kv, total_key_depth) v = (batch, length_kv, total_value_depth)となっていることがわかります。
Args:
q: Tensor with shape [..., length_q, depth_k].
k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
match with q.
v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
match with q.
と書かれています。
まず、1648行目で、
と計算しています。論文中の式(1)の\(QK^T\)の部分です。logits = tf.matmul(q, k, transpose_b=True)
で計算されています。weightsの形状はlogitsと同じです。 ドロップアウトの処理をした後、1667行目のweights = tf.nn.softmax(logits, name="attention_weights")
で、式(1)の計算が完了します。これは、return tf.matmul(weights, v)
Returns:
Tensor with shape [..., length_q, depth_v].
と書いてあります。
以上の計算の途中で得られる[weights]の形状が(..., length_q, length_kv)となっており、
self-attentionの場合はlength_q=length_kvですので、系列長の2乗で必要になるメモリや計算量が増えていくことになります。
投稿者
bluewidz
0
コメント