MNIST 데이터셋을 내려받은 후 이미지를 텐서로 변환합니다.
코드 13-12 데이터셋을 내려받은 후 텐서 변환
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(
root="../chap13/data", train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(
root="../chap13/data", train=False, transform=transform, download=True)
train_loader = DataLoader(
train_dataset, batch_size=100, shuffle=True, num_workers=4, pin_memory=False)
test_loader = DataLoader(
test_dataset, batch_size=100, shuffle=False, num_workers=4)
모델의 네트워크를 생성합니다. 네트워크는 오토인코더처럼 인코더와 디코더로 구성됩니다.
코드 13-13 인코더 네트워크 생성
class Encoder(nn.Module):
def __init__(self, input_dim, hidden_dim, latent_dim):
super(Encoder, self).__init__()
self.input1 = nn.Linear(input_dim, hidden_dim)
self.input2 = nn.Linear(hidden_dim, hidden_dim)
self.mean = nn.Linear(hidden_dim, latent_dim)
self.var = nn.Linear(hidden_dim, latent_dim)
self.LeakyReLU = nn.LeakyReLU(0.2)
self.training = True
def forward(self, x):
h_ = self.LeakyReLU(self.input1(x))
h_ = self.LeakyReLU(self.input2(h_))
mean = self.mean(h_)
log_var = self.var(h_)
return mean, log_var ------ 인코더 네트워크에서 평균과 분산을 반환