def __init__(self):
super(Discriminator, self).__init__()
self.conv_block = nn.Sequential(
# [Conv2d]
# H_out = (H_in + 2 x padding - dilation x (kernel_size -1) -1) / stride + 1
# [-1, 3, 64, 64] -> [-1, 64, 32, 32]
nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
nn.LeakyReLU(0.2, inplace=True),
# [-1, 64, 32, 32] -> [-1, 128, 16, 16]
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2, inplace=True),
# [-1, 128, 16, 16] -> [-1, 256, 8, 8]
nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2, inplace=True),
# [-1, 256, 8, 8] -> [-1, 512, 4, 4]
nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2, inplace=True),
# [-1, 512, 4, 4] -> [-1, 1024, 1, 1]
nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0),
nn.Sigmoid())
utils.weights_init_normal(self)
def __init__(self, latent_dim, classes, channels):
super(Generator, self).__init__()
self.latent_dim = latent_dim
self.channels = channels
self.classes = classes
# image_size = 64
img_shape = (channels, 64, 64)
self.nn_block = nn.Sequential(
# [ConvTranspose2d]
# H_out = (H_in - 1) x stride - 2 x padding + kernel_size + output_padding
# [-1, input_dim, 1, 1] -> [-1, 512, 4, 4]
nn.Linear(latent_dim + classes, 128),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(128, 256),
nn.BatchNorm1d(256, 0.8),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(256, 512),
nn.BatchNorm1d(512, 0.8),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 1024),
nn.BatchNorm1d(1024, 0.8),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(1024, int(np.prod(img_shape))),
nn.Tanh())
utils.weights_init_normal(self)
def __init__(self, latent_dim, channels):
super(Generator, self).__init__()
self.latent_dim = latent_dim
self.channels = channels
self.conv_block = nn.Sequential(
# [ConvTranspose2d]
# H_out = (H_in - 1) x stride - 2 x padding + kernel_size + output_padding
# [-1, input_dim, 1, 1] -> [-1, 512, 4, 4]
nn.ConvTranspose2d(latent_dim, 512, kernel_size=4, stride=1, padding=0),
nn.BatchNorm2d(512),
nn.ReLU(True),
# [-1, 512, 4, 4] -> [-1, 256, 8, 8]
nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(True),
# [-1, 256, 8, 8] -> [-1, 128, 16, 16]
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(True),
# [-1, 128, 16, 16] -> [-1, 64, 32, 32]
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(True),
# [-1, 64, 32, 32] -> [-1, channels, 64, 64]
nn.ConvTranspose2d(64, channels, kernel_size=4, stride=2, padding=1),
nn.Tanh()
)
utils.weights_init_normal(self)
def __init__(self, N_STATES, N_ACTIONS, H1Size, H2Size):
super(Q_Net, self).__init__()
# build network layers
self.fc1 = nn.Linear(N_STATES, H1Size)
self.fc2 = nn.Linear(H1Size, H2Size)
self.out = nn.Linear(H2Size, N_ACTIONS)
# initialize layers
utils.weights_init_normal([self.fc1, self.fc2, self.out], 0.0, 0.1)
def __init__(self, N_STATES, N_ACTIONS, H1Size, H2Size):
super(EnvModel, self).__init__()
# build network layers
self.fc1 = nn.Linear(N_STATES + N_ACTIONS, H1Size)
self.fc2 = nn.Linear(H1Size, H2Size)
self.statePrime = nn.Linear(H2Size, N_STATES)
self.reward = nn.Linear(H2Size, 1)
self.done = nn.Linear(H2Size, 1)
# initialize layers
utils.weights_init_normal(
[self.fc1, self.fc2, self.statePrime, self.reward, self.done], 0.0,
0.1)
def __init__(self, categorical_dim):
super(Discriminator, self).__init__()
self.categorical_dim = categorical_dim
self.conv_block = nn.Sequential(
# [Conv2d]
# H_out = (H_in + 2 x padding - dilation x (kernel_size -1) -1) / strid + 1
# [-1, channels, 64, 64] -> [-1, 64, 32, 32]
nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
nn.LeakyReLU(0.1, inplace=True),
# [-1, 64, 32, 32] -> [-1, 128, 16, 16]
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.1, inplace=True),
# [-1, 128, 16, 16] -> [-1, 256, 8, 8]
nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.1, inplace=True),
# [-1, 256, 8, 8] -> [-1, 512, 4, 4]
nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.1, inplace=True),
# [-1, 512, 4, 4] -> [-1, 64, 1, 1]
# 갑자기 이렇게 체널이 줄어도 괜찮을까?
nn.Conv2d(512, 64, kernel_size=4, stride=1, padding=0),
)
self.adv_layer = nn.Sequential(nn.Linear(64, 1), nn.Sigmoid())
self.aux_layer = nn.Sequential(nn.Linear(64, categorical_dim),
nn.Softmax())
self.latent_layer = nn.Sequential(nn.Linear(64, 1))
utils.weights_init_normal(self)
def weight_init(self, mean=0.0, std=0.02):
for m in self.modules():
utils.weights_init_normal(m, mean=mean, std=std)
netg_b2a = G_net(input_channel, output_channel, ngf, g_layer).to(device)
netd_a = D_net(input_channel, ndf, d_layer).to(device)
netd_b = D_net(input_channel, ndf, d_layer).to(device)
criterionGAN = PatchLoss().to(device)
criterionL1 = nn.L1Loss().to(device)
# criterionMSE = nn.MSELoss().to(device)
optimzer_g = torch.optim.SGD(itertools.chain(netg_b2a.parameters(),
netg_a2b.parameters()),
lr=lr)
optimzerd_a = torch.optim.SGD(netd_a.parameters(), lr)
optimzerd_b = torch.optim.SGD(netd_b.parameters(), lr)
if not os.path.exists(check):
print('init param')
weights_init_normal(optimzer_g)
weights_init_normal(optimzerd_a)
weights_init_normal(optimzerd_b)
def train():
for epoch in range(epochs):
avg_loss_g_a2b = AverageMeter()
avg_loss_g_b2a = AverageMeter()
avg_loss_g = AverageMeter()
avg_loss_d_a = AverageMeter()
avg_loss_d_b = AverageMeter()
min_loss_g = float('inf')
min_loss_d = float('inf')
for i, data in enumerate(train_loader):
img_a, img_b = data[0].to(device), data[1].to(device)
criterionL1 = nn.L1Loss().to(device)
criterionBCE = PatchLoss(nn.BCEWithLogitsLoss()).to(device)
optimzer_g = torch.optim.Adam(itertools.chain(netg_b2a.parameters(),
netg_a2b.parameters()),
lr=lr,
betas=(0.5, 0.999),
weight_decay=weight_decay)
optimzer_d = torch.optim.Adam(itertools.chain(netd_a.parameters(),
netd_b.parameters()),
lr=lr,
betas=(0.5, 0.999),
weight_decay=weight_decay)
if not os.path.exists(check):
print('init param')
weights_init_normal(optimzer_g)
weights_init_normal(optimzer_d)
def train():
for epoch in range(epochs):
avg_loss_g_a = AverageMeter()
avg_loss_g_b = AverageMeter()
avg_loss_d_a = AverageMeter()
avg_loss_d_b = AverageMeter()
min_loss_g = float('inf')
min_loss_d = float('inf')
for i, data in enumerate(train_loader):
img_a, img_b = data[0].to(device), data[1].to(device)
#### update generator
def weight_init(self):
for m in self.modules():
utils.weights_init_normal(m)