代码如下

```

import paddle
from paddle import nn
from paddle import optimizer
from paddle.vision import datasets
from paddle.io import Dataset, DataLoader
from paddle.vision import transforms as T
import numpy as np

transforms = T.Compose([T.ToTensor(), T.Normalize(mean=[127.5], std=[127.5])])

train_data = datasets.MNIST(mode='train', download=True, transform=transforms)
test_data = datasets.MNIST(mode='test', download=True, transform=transforms)

class MyDataset(Dataset):

def __init__(self, train=True):
self.train = train
if train:
self.data = train_data
else:
self.data = test_data

def __getitem__(self, index):
return self.data[index]

def __len__(self):
return len(self.data)

BATCH_SIZE = 128
train_loader = DataLoader(MyDataset(train=True), batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(MyDataset(train=False), batch_size=BATCH_SIZE, shuffle=False)
net = nn.Sequential(
nn.Flatten(),
nn.Linear(784, 512),
nn.ReLU(),
nn.Linear(512, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 10)
)
net.train()
optim = optimizer.Adam( parameters=net.parameters())
criterion = paddle.nn.CrossEntropyLoss()

for i, (x, y) in enumerate(train_loader):
y = paddle.to_tensor(y, stop_gradient=False)
y_pred = net(x)
loss = criterion(y_pred, y)
loss.backward()
optim.step()
optim.clear_grad()
if (i + 1) % 50 == 0:
print(f'loss: {loss.numpy()[0]:.2f}')

```