Test set accuracy is very high after very few epochs on mnist dataset
With very few epochs this model learns to classify beween 1 and 0 extremely quickly which leads me to consider something is wrong. 

Below code downloads mnist dataset, extracts the mnist images that contain 1 or 0 only. A random sample of size 200 is selected from this subset of mnist images. This random sample is the dataset the model is trained on. With just 2 epochs the model achieves 90%+ test set accuracy, is this expected behaviour ? I expected many more epochs would be required in order to train the model to achieve this level of test set accuracy. 

Model code : 

%reset -f

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import torch.utils.data as data_utils
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
from matplotlib import pyplot
from pandas import DataFrame
import torchvision.datasets as dset
import os
import torch.nn.functional as F
import time
import random
import pickle
from sklearn.metrics import confusion_matrix
import pandas as pd
import sklearn


trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])

root = './data'
if not os.path.exists(root):
    os.mkdir(root)
train_set = dset.MNIST(root=root, train=True, transform=trans, download=True)
test_set = dset.MNIST(root=root, train=False, transform=trans, download=True)

batch_size = 64

train_loader = torch.utils.data.DataLoader(
                 dataset=train_set,
                 batch_size=batch_size,
                 shuffle=True)
test_loader = torch.utils.data.DataLoader(
                dataset=test_set,
                batch_size=batch_size,
shuffle=True)

class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(28*28, 500)
        self.fc2 = nn.Linear(500, 256)
        self.fc3 = nn.Linear(256, 2)
    def forward(self, x):
        x = x.view(-1, 28*28)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

num_epochs = 2
random_sample_size = 200

values_0_or_1 = [t for t in train_set if (int(t[1]) == 0 or int(t[1]) == 1)]
values_0_or_1_testset = [t for t in test_set if (int(t[1]) == 0 or int(t[1]) == 1)]

print(len(values_0_or_1))
print(len(values_0_or_1_testset))

train_loader_subset = torch.utils.data.DataLoader(
                 dataset=values_0_or_1,
                 batch_size=batch_size,
                 shuffle=True)

test_loader_subset = torch.utils.data.DataLoader(
                 dataset=values_0_or_1_testset,
                 batch_size=batch_size,
                 shuffle=False)

train_loader = train_loader_subset

# Hyper-parameters 
input_size = 100
hidden_size = 100
num_classes = 2
# learning_rate = 0.00001
learning_rate = .0001
# Device configuration
device = 'cpu'
print_progress_every_n_epochs = 1

model = NeuralNet().to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  

N = len(train_loader)
# Train the model
total_step = len(train_loader)

most_recent_prediction = []
test_actual_predicted_dict = {}

rm = random.sample(list(values_0_or_1), random_sample_size)
train_loader_subset = data_utils.DataLoader(rm, batch_size=4)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader_subset):  
        # Move tensors to the configured device
        images = images.reshape(-1, 2).to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    if (epoch) % print_progress_every_n_epochs == 0:
        print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, i+1, total_step, loss.item()))


predicted_test = []
model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
probs_l = []

predicted_values = []
actual_values = []
labels_l = []

with torch.no_grad():
    for images, labels in test_loader_subset:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        predicted_test.append(predicted.cpu().numpy())

        sm = torch.nn.Softmax()
        probabilities = sm(outputs) 
        probs_l.append(probabilities)  
        labels_l.append(labels.cpu().numpy())

    predicted_values.append(np.concatenate(predicted_test).ravel())
    actual_values.append(np.concatenate(labels_l).ravel())

if (epoch) % 1 == 0:
    print('test accuracy : ', 100 * len((np.where(np.array(predicted_values[0])==(np.array(actual_values[0])))[0])) / len(actual_values[0]))


Output of model (12665 &amp; 2115 represents the training and test set sizes) : 

12665
2115
Epoch [1/2], Step [50/198], Loss: 0.1256
Epoch [2/2], Step [50/198], Loss: 0.0151
test accuracy :  99.76359338061465

/anaconda3/envs/pytorch/lib/python3.7/site-packages/ipykernel_launcher.py:143: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.


Here's my 2 cents on your binary experiment. 

It would seem like you have severely reduce the complexity of your dataset and with the high number of neurons in your intermediate layers, your model is expected to converge very quickly.

Note that MNIST dataset has channel of 1 and this makes the task very simple.

You may try to play with CIFAR10 and see if you are still getting high accuracy in just 2 epochs.
