Define Autoencoder
#dropout_p= 0.5
class Autoencoder(nn.Module):
def __init__(self):
super(Autoencoder, self).__init__()
self.encoder = nn.Sequential(
# nn.Dropout2d(dropout_p= 0.5),
nn.Conv2d(1, 16,kernel_size=3,stride =1, padding= 1,bias= False),
nn.ReLU(True),
nn.MaxPool2d( 2, 2,return_indices=True),
nn.Conv2d(16,32,kernel_size=3,stride =1, padding=1),
nn.ReLU(True),
nn.MaxPool2d(2, 2,return_indices=True),
nn.Conv2d(32,64,kernel_size=3,stride =1, padding= 1),
nn.ReLU(True),
nn.MaxPool2d(2, 2,return_indices=True),
)
self.decoder = nn.Sequential(
nn.MaxUnpool2d(2, 2),
nn.ConvTranspose2d(64, 32, kernel_size=3),
nn.ReLU(True),
nn.MaxUnpool2d(2, 2),
nn.ConvTranspose2d(32, 16, kernel_size=3),
nn.ReLU(True),
nn.MaxUnpool2d(2, 2),
nn.ConvTranspose2d(16, 8, kernel_size=3)
#nn.Tanh()
)
def forward(self, x):
x = self.encoder(x)
encoded_x = x
x = self.decoder(x)
#x=F.dropout(x,training=self.training)
return x, encoded_x
num_epochs = 5
batch_size = 128
learning_rate = 1e-5
model= Autoencoder()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(num_epochs):
for data in loader_train_set:
img, label = data
img = img.view(-1,1,320,240)
img = Variable(img)
# img = torch.from_numpy(numpy.array(img))
encoded_output = model.encoder(img)
loss = criterion(encoded_output, img)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(‘epoch [{}/{}], loss:{:.4f}’.format(epoch+1, num_epochs, loss.data()))