import torch
import numpy as np
import argparse
import torch
import torchvision
import deepspeed

# 命令行参数：deepspeed ds_train.py --epoch 5 --deepspeed --deepspeed_config ds_config.json
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=1, help="local rank passed from distributed launcher")
parser.add_argument("--epoch", type=int, default=1, help="epoch")
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()  # deepspeed命令行参数

class FashionModel(torch.nn.Module): 
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Linear(in_features=784, out_features=300),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=300, out_features=10)
        )

    def forward(self, batch_x):
        return self.seq(batch_x)

def img_transform(img): 
    img = np.asarray(img) / 255
    return torch.tensor(img, dtype=torch.float32).flatten()


dataset = torchvision.datasets.FashionMNIST(root='./dataset', download=True, transform=img_transform)  # 数据集
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=2, shuffle=True)  # 数据加载器，batch_size应该等于train_batch_size/gpu数量

model = FashionModel()  # 自定义的模型
model, _, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=model.parameters())  # deepspeed分布式模型
loss_fn = torch.nn.CrossEntropyLoss()

for epoch in range(cmd_args.epoch):
    for x, y in dataloader:
        x = x.cuda()
        y = y.cuda()
        output = model(x)
        loss = loss_fn(output, y)
        model.backward(loss)  # 走deepspeed风格的backward
        model.step()
    print("epoch {} done...".format(epoch))
    model.save_checkpoint('./checkpoints')