Repository navigation
Expand file tree
/
Copy pathTrainProcess.py
More file actions
114 lines (86 loc) · 4.24 KB
/
Copy pathTrainProcess.py
File metadata and controls
114 lines (86 loc) · 4.24 KB
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
from torch.nn.modules.loss import BCEWithLogitsLoss
from Dataset import ModelNet10
from GAN import GAN, compute_norm
from torch.optim import AdamW
from torch.utils.tensorboard import SummaryWriter
from Generator import Generator
from VersionManager import VersionManager
import torch
import numpy as np
class TrainProcesss:
def __init__(self, model: GAN, vm: VersionManager, epochs=1000, device=torch.device('cpu')) -> None:
torch.autograd.set_detect_anomaly(True)
self.gan = model
self.ds = ModelNet10()
self.vm = vm
self.writer = SummaryWriter(f'runs/{self.vm.name}')
self.cepoch = 0
self.device = device
self.epochs = epochs
self.gstep = 0
self.doptim = AdamW(self.gan.discriminator.parameters(), lr=1e-4, betas=(0.5, 0.999))
self.goptim = AdamW(self.gan.generator.parameters(), lr=1e-4, betas=(0.5, 0.999))
self.dscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
self.doptim, T_max=self.epochs, eta_min=1e-8
)
self.gscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
self.goptim, T_max=self.epochs, eta_min=1e-7
)
self.lossess = []
self.criterion = BCEWithLogitsLoss()
def mbgd_step(self, step_fn):
avgd_loss = []
avgg_loss = []
mbgd_epoch = 1
for x_batch in self.ds.loader:
g, d = step_fn(x_batch)
avgg_loss.append(g.detach().numpy())
avgd_loss.append(d.detach().numpy())
mbgd_epoch += 1
self.gstep += 1
return np.mean(avgg_loss), np.mean(avgd_loss)
def step_fn(self, x_batch):
self.doptim.zero_grad()
self.goptim.zero_grad()
real = x_batch.x
real = real.detach().unsqueeze(1)
latent_vector = torch.randn(self.ds.batch_size, 100, 1, 1, 1, device=self.device)
fake = self.gan.generator(latent_vector)
ddl = self.gan.compute_detailed_d_loss(self.criterion, real, fake)
d_loss, dreal, dfake, preal, pfake = ddl
acc = ((preal > 0).float().mean() + (pfake < 0).float().mean())/2
# D is updated only when its accuracy < 80%, following the 3D-GAN training scheme
# (Wu et al., NeurIPS 2016, https://arxiv.org/abs/1610.07584).
if acc < 0.8:
self.doptim.zero_grad()
d_loss.backward()
self.doptim.step()
self.doptim.zero_grad()
self.goptim.zero_grad()
g_loss = self.gan.compute_g_loss(self.criterion, fake)
g_loss.backward()
gnorm = compute_norm(self.gan.generator)
self.goptim.step()
self.writer.add_scalar("Batch/Generator/Loss", g_loss.item(), self.gstep)
self.writer.add_scalar("Batch/Generator/Gradient Norm", gnorm.item(), self.gstep)
self.writer.add_scalar("Batch/Discriminator/Loss", d_loss.item(), self.gstep)
self.writer.add_scalar("Batch/Discriminator/Real Loss", dreal.item(), self.gstep)
self.writer.add_scalar("Batch/Discriminator/Fake Loss", dfake.item(), self.gstep)
self.writer.add_scalar("Batch/Discriminator/Real Score", preal.detach().numpy().mean(), self.gstep)
self.writer.add_scalar("Batch/Discriminator/Fake Score", pfake.detach().numpy().mean(), self.gstep)
self.writer.add_scalars('Batch/Loss Comparison', {'generator': g_loss.item(), 'discriminator': d_loss.item()}, self.gstep)
return (g_loss, d_loss)
def train(self):
self.gan.discriminator.train()
self.gan.generator.train()
for e in range(self.epochs):
self.vm.setEpoch(self.cepoch)
ag, ad = self.mbgd_step(self.step_fn)
self.writer.add_scalar("Epoch/Generator/Loss", ag.item(), self.cepoch)
self.writer.add_scalar("Epoch/Discriminator/Loss", ad.item(), self.cepoch)
self.writer.add_scalars('Epoch/Loss Comparison', {'generator': ag.item(), 'discriminator': ad.item()}, self.cepoch)
self.cepoch += 1
self.lossess.append((ag, ad))
self.dscheduler.step()
self.gscheduler.step()
self.vm.save()