Commit @415bdf44a549e535f776f2c23546a6f76ff563e5

윤영준 2023-06-27

Discriminator Loss

@415bdf44a549e535f776f2c23546a6f76ff563e5

5cda0f7

415bdf4

model/Discriminator.py

--- model/Discriminator.py

+++ model/Discriminator.py


         self.fc1 = nn.Linear(32, 1)  # You need to adjust the input dimension here depending on your input size
         self.fc2 = nn.Linear(1, 1)
 
+    def loss(self, input_tensor, label_tensor, attention_map, name):
+        batch_size, image_h, image_w, _ = input_tensor.size()
+
+        zeros_mask = torch.zeros([batch_size, image_h, image_w, 1], dtype=torch.float32)
+
+        # Inference function
+        fc_out_o, attention_mask_o, fc2_o = self.forward(input_tensor)
+        fc_out_r, attention_mask_r, fc2_r = self.forward(label_tensor)
+
+        l_map = F.mse_loss(attention_map, attention_mask_o) + \
+                F.mse_loss(attention_mask_r, zeros_mask)
+
+        entropy_loss = -torch.log(fc_out_r) - torch.log(-torch.sub(fc_out_o, 1.0))
+        entropy_loss = torch.mean(entropy_loss)
+
+        loss = entropy_loss + 0.05 * l_map
+
+        return fc_out_o, loss
+
     def forward(self, x):
         x1 = F.leaky_relu(self.conv1(x))
         x2 = F.leaky_relu(self.conv2(x1))

5cda0f7

415bdf4

tools/dataloader.py

--- tools/dataloader.py

+++ tools/dataloader.py


         self.clean_img = clean_img_dir
         self.rainy_img = rainy_img_dirs
         self.transform = transform
-        # self.target_transform = target_transform
-        # print(self.transform)
 
     def __len__(self):
         return len(self.clean_img)

5cda0f7

415bdf4

train.py

--- train.py

+++ train.py


     generator.apply(weights_init_normal)
     discriminator.apply(weights_init_normal)
 
-dataloader = Dataloader
+dataloader = Dataloader()
 
 
+# declare generator loss
+
+
+for epoch in range(epoch):
+    for i, (imgs, _) in enumerate(dataloader):
+
+
+torch.save(generator.attentionRNN.state_dict(), "attentionRNN_model_path")
+
 
 
 ## RNN 따로 돌리고 CPU로 메모리 옳기고

Add a comment

Open 0
Closed 0

List

...	...	@@ -17,6 +17,25 @@
17	17	self.fc1 = nn.Linear(32, 1) # You need to adjust the input dimension here depending on your input size
18	18	self.fc2 = nn.Linear(1, 1)
19	19
	20	+ def loss(self, input_tensor, label_tensor, attention_map, name):
	21	+ batch_size, image_h, image_w, _ = input_tensor.size()
	22	+
	23	+ zeros_mask = torch.zeros([batch_size, image_h, image_w, 1], dtype=torch.float32)
	24	+
	25	+ # Inference function
	26	+ fc_out_o, attention_mask_o, fc2_o = self.forward(input_tensor)
	27	+ fc_out_r, attention_mask_r, fc2_r = self.forward(label_tensor)
	28	+
	29	+ l_map = F.mse_loss(attention_map, attention_mask_o) + \
	30	+ F.mse_loss(attention_mask_r, zeros_mask)
	31	+
	32	+ entropy_loss = -torch.log(fc_out_r) - torch.log(-torch.sub(fc_out_o, 1.0))
	33	+ entropy_loss = torch.mean(entropy_loss)
	34	+
	35	+ loss = entropy_loss + 0.05 * l_map
	36	+
	37	+ return fc_out_o, loss
	38	+
20	39	def forward(self, x):
21	40	x1 = F.leaky_relu(self.conv1(x))
22	41	x2 = F.leaky_relu(self.conv2(x1))

...	...	@@ -8,8 +8,6 @@
8	8	self.clean_img = clean_img_dir
9	9	self.rainy_img = rainy_img_dirs
10	10	self.transform = transform
11		- # self.target_transform = target_transform
12		- # print(self.transform)
13	11
14	12	def __len__(self):
15	13	return len(self.clean_img)

...	...	@@ -44,9 +44,18 @@
44	44	generator.apply(weights_init_normal)
45	45	discriminator.apply(weights_init_normal)
46	46
47		-dataloader = Dataloader
	47	+dataloader = Dataloader()
48	48
49	49
	50	+# declare generator loss
	51	+
	52	+
	53	+for epoch in range(epoch):
	54	+ for i, (imgs, _) in enumerate(dataloader):
	55	+
	56	+
	57	+torch.save(generator.attentionRNN.state_dict(), "attentionRNN_model_path")
	58	+
50	59
51	60
52	61	## RNN 따로 돌리고 CPU로 메모리 옳기고

Delete comment