Pytorch随笔

代码链接：https://github.com/zhuqunxi/pytorch-implement-NLP

P01 -- Two layer model

1. Numpy to tensor: x_tensor = torch.from_numpy(np_x)
2. Cpu tensor to cuda: x_tensor_cuda= x_tensor.cuda()
3. Cuda data to Variable: x_tensor_cuda_var=Variable(x_tensor_cuda)
4. Tensor to numpy: x_np=x_tensor.cpu().numpy()
5. Variable to numpy: x_np=x_tensor_cuda_var.cpu().detach().numpy()

 　　随机数据

import numpy as np
import torch
import torch.nn as nn
np.random.seed(1)
torch.manual_seed(1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
print('device:', device)
device = torch.device(device)

N, D_in, D_out =64, 1000, 10
train_x = np.random.normal(size=(N, D_in))
train_y = np.random.normal(size=(N, D_out))

class Two_layer(torch.nn.Module):
    def __init__(self, D_in, D_out, H=100):
        super(Two_layer, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(H, D_out)
    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

model = Two_layer(D_in, D_out, H=1000)

loss_fn = nn.MSELoss(reduction='sum')
# optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-4)
optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-4)

train_x = torch.from_numpy(train_x).type(dtype=torch.float32)
train_y = torch.from_numpy(train_y).type(dtype=torch.float32)

# train_x = torch.randn(N, D_in)
# train_y = torch.randn(N, D_out)

train_x = train_x.to(device)
train_y = train_y.to(device)
model = model.to(device)

import time
time_st = time.time()
for epoch in range(200):
    y_pred = model(train_x)

    loss = loss_fn(y_pred, train_y)
    optimizer.zero_grad()
    if not epoch % 20:
        print('loss: ', loss.item())
    loss.backward()
    optimizer.step()
print('training time used {:.1f} s with {}'.format(time.time() - time_st, device))

"""
loss:  673.6837158203125
loss:  57.70276641845703
loss:  3.7402660846710205
loss:  0.2832883596420288
loss:  0.026732178404927254
loss:  0.0029198969714343548
loss:  0.00034921077894978225
loss:  4.434480797499418e-05
loss:  5.87546583119547e-06
loss:  8.037222301027214e-07
training time used 1.1 s with cpu
training time used 0.6 s with cuda
"""

 　　Mnist数据集

from keras.datasets import mnist
import torch
import numpy as np
np.random.seed(1)
torch.manual_seed(1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
print('device:', device)
device = torch.device(device)

(x_train, y_train), (x_test, y_test) = mnist.load_data()
print('x_train, y_train shape:', x_train.shape, y_train.shape)
print('x_test, y_test shape:', x_test.shape, y_test.shape)
x_train = x_train.reshape(x_train.shape[0], -1)
x_test = x_test.reshape(x_test.shape[0], -1)
print('x_train, y_train shape:', x_train.shape, y_train.shape)
print('x_test, y_test shape:', x_test.shape, y_test.shape)

N, D_in, D_out = 1000, x_train.shape[1], 10

class Two_layer(torch.nn.Module):
    def __init__(self, D_in, D_out, H=1000):
        super(Two_layer, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(H, D_out)
    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

model = Two_layer(D_in, D_out, H = 1000)
loss_fn = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(params=model.parameters(), lr=1e-4)
x_train, y_train = torch.tensor(x_train,dtype=torch.float32), torch.tensor(y_train, dtype=torch.long)
x_test, y_test = torch.tensor(x_test,dtype=torch.float32), torch.tensor(y_test, dtype=torch.long)

x_train, y_train = x_train.to(device), y_train.to(device)
x_test, y_test = x_test.to(device), y_test.to(device)
model = model.to(device)
import time
time_st = time.time()
for epoch in range(50):
    y_pred = model(x_train)
    loss = loss_fn(y_pred, y_train)

    if not epoch % 10:
        with torch.no_grad():
            y_pred_test = model(x_test)
            y_label_pred = np.argmax(y_pred_test.cpu().detach().numpy(), axis=1)
            print('y_label_pred y_test shape:', y_label_pred.shape, y_test.size())
            acc_test = np.mean(y_label_pred == y_test.cpu().detach().numpy())
            loss_test = loss_fn(y_pred_test, y_test)
            print('test loss: {}, acc: {}'.format(loss_test.item(), acc_test))

            y_label_pred_train = np.argmax(y_pred.cpu().detach().numpy(), axis=1)
            acc_train = np.mean(y_label_pred_train == y_train.cpu().detach().numpy())
            print('train loss: {}, acc: {}'.format(loss.item(), acc_train))

            print('-' * 80)

    opt.zero_grad()
    loss.backward()
    opt.step()

print('training time used {:.2f} s with device {}'.format(time.time() - time_st, device))

'''
x_train, y_train shape: (60000, 28, 28) (60000,)
x_test, y_test shape: (10000, 28, 28) (10000,)
x_train, y_train shape: (60000, 784) (60000,)
x_test, y_test shape: (10000, 784) (10000,)
y_label_pred y_test shape: (10000,) torch.Size([10000])
test loss: 23.847854614257812, acc: 0.1414
train loss: 23.87252426147461, acc: 0.13683333333333333
--------------------------------------------------------------------------------
y_label_pred y_test shape: (10000,) torch.Size([10000])
test loss: 3.340665578842163, acc: 0.7039
train loss: 3.514056444168091, acc: 0.6925166666666667
--------------------------------------------------------------------------------
y_label_pred y_test shape: (10000,) torch.Size([10000])
test loss: 1.7213207483291626, acc: 0.844
train loss: 1.8277908563613892, acc: 0.84025
--------------------------------------------------------------------------------
y_label_pred y_test shape: (10000,) torch.Size([10000])
test loss: 1.2859240770339966, acc: 0.8845
train loss: 1.3402273654937744, acc: 0.88125
--------------------------------------------------------------------------------
y_label_pred y_test shape: (10000,) torch.Size([10000])
test loss: 1.0803418159484863, acc: 0.8993
train loss: 1.084514856338501, acc: 0.8984833333333333
--------------------------------------------------------------------------------
training time used 81.26 s with device cpu
training time used 3.61 s with device cuda
'''

P02 wordvec

　　Skipgram model

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import os

np.random.seed(1)
torch.manual_seed(1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
print('device:', device)
device = torch.device(device)

small = 5000000
training = True
N_epoch = 51
Batch_size = 128
show_epoch = 1

words = {}
with open('zhihu.txt', mode='r', encoding='utf8') as f:
    lines = f.readlines()
    print('len(lines):', len(lines))
    for idx, line in enumerate(lines):
        # print(line)
        for word in line.split():
            if word in words:
                words[word] += 1
            else:
                words[word] = 1
        if idx > small:
            break
print(len(words))
# print(words)
word2index = {word: idx for idx, word in enumerate(words.keys())}
indx2word = {idx: word for idx, word in enumerate(words.keys())}
# print(word2index)
# print(indx2word)
word_freq = np.array(list(words.values()))
word_freq = word_freq / np.sum(word_freq)
word_freq = word_freq ** (3 / 4.0)
word_freq = word_freq / np.sum(word_freq)
word_freq = torch.Tensor(word_freq)
# print(word_freq)

C, K = 3, 10 # C:窗口大小， K:每个positive样本对应K个negative样本
em_dim = 100
word_size = len(words)


def creat_train_data():
    Center_Outside_words, Center_Outside_words_index = [], []
    with open('zhihu.txt', mode='r', encoding='utf8') as f:
        lines = f.readlines()
        print('len(lines):', len(lines))
        for _, line in enumerate(lines):
            # print(line)
            line = line.split()
            n = len(line)
            for idx, word in enumerate(line):
                st = max(idx - C, 0)
                ed = min(idx + 1 + C, n)
                for i in range(st, idx):
                    word_ = line[i]
                    Center_Outside_words.append([word, word_])
                    Center_Outside_words_index.append([word2index[word], word2index[word_]])
                for i in range(idx + 1, ed):
                    word_ = line[i]
                    Center_Outside_words.append([word, word_])
                    Center_Outside_words_index.append([word2index[word], word2index[word_]])
            if _ > small:
                break
    return Center_Outside_words, Center_Outside_words_index

Center_Outside_words, Center_Outside_words_index = creat_train_data()
Center_Outside_words_index = np.array(Center_Outside_words_index)

print(Center_Outside_words[:10])
print(Center_Outside_words_index[:10])
print('train data len:', len(Center_Outside_words))

N_train = len(Center_Outside_words)

def get_batch(batch_step):
    st, ed = batch_step * Batch_size, min(batch_step * Batch_size + Batch_size, N_train)
    assert st < ed
    center_word = torch.LongTensor(Center_Outside_words_index[st:ed, 0])  # (batch, )
    outside_word = torch.LongTensor(Center_Outside_words_index[st:ed, 1]) # (batch, )
    negtive_word = torch.multinomial(word_freq, K * (ed - st)).view(-1, K) # (batch, K)

    # print(center_word.size(), outside_word.size(), negtive_word.size())
    # print(center_word, outside_word, negtive_word)
    return center_word, outside_word, negtive_word

center_word, outside_word, negtive_word = get_batch(batch_step=0)


class Zhihu_DataSet(Dataset):
    def __init__(self, Center_Outside_words_index, word_freq):
        self.Center_Outside_words_index = Center_Outside_words_index
        self.word_freq = word_freq
        print('Center_Outside_words_index shape:', Center_Outside_words_index.shape)

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

    def __getitem__(self, index):
        # center_word = torch.LongTensor([self.Center_Outside_words_index[index, 0]])
        # outside_word = torch.LongTensor([self.Center_Outside_words_index[index, 1]])

        center_word = torch.tensor(self.Center_Outside_words_index[index, 0],dtype=torch.long)
        outside_word = torch.tensor(self.Center_Outside_words_index[index, 1],dtype=torch.long)

        negtive_word = torch.multinomial(word_freq, K, replacement=True)  # (batch, K)
        # print(center_word.size(), outside_word.size(), negtive_word.size())
        return center_word, outside_word, negtive_word



zhihu_dataset = Zhihu_DataSet(Center_Outside_words_index, word_freq)
zhihu_dataloader = DataLoader(dataset=zhihu_dataset,batch_size=Batch_size, shuffle=True)

class Word2Vec_Zqx(nn.Module):
    def __init__(self, word_size, em_dim):
        super(Word2Vec_Zqx, self).__init__()
        self.word_em_center = nn.Embedding(num_embeddings=word_size,embedding_dim=em_dim)
        self.word_em_outside = nn.Embedding(num_embeddings=word_size,embedding_dim=em_dim)

    def forward(self, center_word, outside_word, negtive_word):
        center_word_emd = self.word_em_center(center_word)  # (batch, em_dim)
        outside_word_emd = self.word_em_outside(outside_word) # (batch, em_dim)
        negtive_word_emd = self.word_em_outside(negtive_word) # (batch, K, em_dim))

        # print(center_word_emd.size(), outside_word_emd.size(), negtive_word_emd.size())
        center_word_emd = center_word_emd.unsqueeze(dim=2)  # (batch, em_dim, 1)
        outside_word_emd = outside_word_emd.unsqueeze(dim=1)  # (batch, 1, em_dim)
        # print(center_word_emd.size(), outside_word_emd.size(), negtive_word_emd.size())
        center_outside_word = torch.bmm(outside_word_emd, center_word_emd).squeeze(1)
        center_outside_word = center_outside_word.squeeze(1)  # (batch, )
        center_negtive_word = torch.bmm(negtive_word_emd, center_word_emd).squeeze(2)  # (batch, K)
        # print(center_outside_word.size(), center_negtive_word.size())

        loss = - (torch.sum(F.logsigmoid(center_outside_word)) + torch.sum(F.logsigmoid(center_negtive_word)))
        return loss

    def get_emd_center(self):
        return self.word_em_center.weight.cpu().detach().numpy()

model =Word2Vec_Zqx(word_size=word_size, em_dim=em_dim)
loss = model(center_word, outside_word, negtive_word)
print('loss:', loss.item())

# 模型保存
check_path = './Checkpoints/'
filepath = check_path + 'word2vec_state_dict.pkl'
def find_similar_word(emd_center, word):
    word_idx = word2index[word]
    word_emd = emd_center[word_idx].reshape(-1, 1)
    # similarity = np.matmul(emd_center, word_emd).flatten()
    similarity = np.matmul(emd_center, word_emd).flatten() / np.linalg.norm(emd_center, axis=1) / np.linalg.norm(word_emd)
    k = 10
    topk_idx = np.argsort(-similarity)[:k]

    print('与word=[{}]--相似的top {}的有:'.format(word, k))
    topk_word = [indx2word[_] for _ in topk_idx]
    print(topk_word)

def train(model):
    # opt = torch.optim.Adam(model.parameters(), lr=1e-4)
    opt = torch.optim.SGD(model.parameters(), lr=1e-2)
    model.to(device)
    import time
    time_st_global = time.time()
    for epoch in range(N_epoch):
        time_st_epoch = time.time()
        for batch_step in range(N_train // Batch_size):
            center_word, outside_word, negtive_word = get_batch(batch_step)
            center_word, outside_word, negtive_word = center_word.to(device), outside_word.to(device), negtive_word.to(device)
            loss = model(center_word, outside_word, negtive_word)

            opt.zero_grad()
            loss.backward()
            opt.step()
        print('# ' * 80)
        print('epoch:{}, batch_step: {}, loss: {}'.format(epoch, batch_step, loss.item()))
        print('epoch time {:.2f} s, total time {:.2f} s'.format(time.time() - time_st_epoch, time.time() - time_st_global))
        if not epoch % show_epoch:
            if not os.path.exists(check_path):
                os.makedirs(check_path)
            torch.save(model.state_dict(), filepath)
            emd_center = model.get_emd_center()

            test_words = ['你', '为什么', '学生', '女生', '什么', '大学']
            for word in test_words:
                print('-' * 80)
                print('test word : {},  次数: {}'.format(word, words[word]))
                find_similar_word(emd_center=emd_center, word=word)

    return model

def train_with_dataloader(model):
    # opt = torch.optim.Adam(model.parameters(), lr=1e-4)
    opt = torch.optim.SGD(model.parameters(), lr=1e-2)
    model.to(device)
    import time
    time_st_global = time.time()
    for epoch in range(N_epoch):
        time_st_epoch = time.time()
        for batch_step, (center_word, outside_word, negtive_word) in enumerate(zhihu_dataloader):
            center_word, outside_word, negtive_word = center_word.to(device), outside_word.to(device), negtive_word.to(device)
            loss = model(center_word, outside_word, negtive_word)

            opt.zero_grad()
            loss.backward()
            opt.step()
        print('#' * 80)
        print('epoch:{}, batch_step: {}, loss: {}'.format(epoch, batch_step, loss.item()))
        print('epoch time {:.2f} s, total time {:.2f} s'.format(time.time() - time_st_epoch, time.time() - time_st_global))
        if not epoch % show_epoch:
            if not os.path.exists(check_path):
                os.makedirs(check_path)
            torch.save(model.state_dict(), filepath)

            # emd_center = model.get_emd_center()
            # test_words = ['你', '为什么', '学生', '女生', '什么', '大学']
            # for word in test_words:
            #     print('-' * 80)
            #     print('test word : {},  次数: {}'.format(word, words[word]))
            #     find_similar_word(emd_center=emd_center, word=word)

    return model

if training:
    # model = train(model)
    model = train_with_dataloader(model)

# 模型恢复
model.load_state_dict(torch.load(filepath))

emd_center = model.get_emd_center()


test_words = ['你', '为什么', '学生', '女生', '什么', '大学']

for word in test_words:
    print('-' * 80)
    print('test word : {},  次数: {}'.format(word, words[word]))
    find_similar_word(emd_center=emd_center, word=word)

print('end!!!')

P03 RNN

　　mnist classification

from keras.datasets import mnist
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time
np.random.seed(1)
torch.manual_seed(1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
print('device:', device)
device = torch.device(device)

N, D_in, D_out = 10000, 28, 10
H = 100
Batch_size = 128
lr=1e-2
N_epoch = 200


(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, y_train = x_train[:N], y_train[:N]
x_test, y_test = x_test[:N], y_test[:N]

# 归一化很重要，不然有可能train不起来，或者test效果不行
x_train = x_train /255.0
x_test = x_test / 255.0

print('x_train, y_train shape:', x_train.shape, y_train.shape)
print('x_test, y_test shape:', x_test.shape, y_test.shape)
print('np.max(x_train), np.min(x_train):', np.max(x_train), np.min(x_train))
print('np.max(y_train), np.min(y_train):', np.max(y_train), np.min(y_train))

class RNN_zqx(nn.Module):
    def __init__(self, D_in, H):
        super(RNN_zqx, self).__init__()
        self.rnn = nn.LSTM(input_size=D_in,hidden_size=H,num_layers=1,batch_first=True)
        self.linear = nn.Linear(H, 10)
    def forward(self, x):
        all_h, (h, c) = self.rnn(x)
        # all_h: (batch, seq_len, num_directions * hidden_size)
        # h: (num_layers * num_directions, batch, hidden_size)
        # print('all_h.size():', all_h.size())
        # print('h.size():', h.size())
        x = self.linear(h.squeeze(0))
        return x

model =RNN_zqx(D_in=D_in, H=H)
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=lr)

x_train, y_train = torch.Tensor(x_train), torch.LongTensor(y_train)
x_test, y_test = torch.Tensor(x_test), torch.LongTensor(y_test)

print('x_train.size(), y_train.size():', x_train.size(), y_train.size())
x_train, y_train = x_train.to(device), y_train.to(device)
x_test, y_test = x_test.to(device), y_test.to(device)
mdoel = model.to(device)

time_st = time.time()
for epoch in range(N_epoch):
    y_pred = model(x_train)
    # print(y_pred.size())
    loss = loss_fn(y_pred, y_train)

    if not epoch % 10:
        with torch.no_grad():
            y_pred_test = model(x_test)
            y_label_pred = np.argmax(y_pred_test.cpu().detach().numpy(), axis=1)
            # print('y_label_pred y_test shape:', y_label_pred.shape, y_test.size())
            acc_test = np.mean(y_label_pred == y_test.cpu().detach().numpy())
            loss_test = loss_fn(y_pred_test, y_test)
            print('test loss: {}, acc: {}'.format(loss_test.item(), acc_test))

            y_label_pred_train = np.argmax(y_pred.cpu().detach().numpy(), axis=1)
            acc_train = np.mean(y_label_pred_train == y_train.cpu().detach().numpy())
            print('train loss: {}, acc: {}'.format(loss.item(), acc_train))

            print('-' * 80)

    opt.zero_grad()
    loss.backward()
    opt.step()

print('Training time used {:.2f} s'.format(time.time() - time_st))

'''
device: cuda
x_train, y_train shape: (10000, 28, 28) (10000,)
x_test, y_test shape: (10000, 28, 28) (10000,)
np.max(x_train), np.min(x_train): 1.0 0.0
np.max(y_train), np.min(y_train): 9 0
x_train.size(), y_train.size(): torch.Size([10000, 28, 28]) torch.Size([10000])
test loss: 2.3056862354278564, acc: 0.1032
train loss: 2.3057758808135986, acc: 0.0991
--------------------------------------------------------------------------------
test loss: 1.6542853116989136, acc: 0.5035
train loss: 1.651445746421814, acc: 0.482
--------------------------------------------------------------------------------
test loss: 1.0779469013214111, acc: 0.6027
train loss: 1.0364742279052734, acc: 0.6158
--------------------------------------------------------------------------------
test loss: 0.7418596148490906, acc: 0.7503
train loss: 0.7045448422431946, acc: 0.7642
--------------------------------------------------------------------------------
test loss: 0.5074136853218079, acc: 0.8369
train loss: 0.46816474199295044, acc: 0.8512
--------------------------------------------------------------------------------
test loss: 0.3507310748100281, acc: 0.8931
train loss: 0.29413318634033203, acc: 0.9125
--------------------------------------------------------------------------------
test loss: 0.25384169816970825, acc: 0.9292
train loss: 0.1905861645936966, acc: 0.9446
--------------------------------------------------------------------------------
test loss: 0.21215158700942993, acc: 0.9406
train loss: 0.13411203026771545, acc: 0.9614
--------------------------------------------------------------------------------
test loss: 0.19598548114299774, acc: 0.9467
train loss: 0.0968935638666153, acc: 0.9711
--------------------------------------------------------------------------------
test loss: 0.6670947074890137, acc: 0.834
train loss: 0.6392199993133545, acc: 0.8405
--------------------------------------------------------------------------------
test loss: 0.3550219237804413, acc: 0.8966
train loss: 0.29769250750541687, acc: 0.9112
--------------------------------------------------------------------------------
test loss: 0.22847041487693787, acc: 0.9345
train loss: 0.16787868738174438, acc: 0.9545
--------------------------------------------------------------------------------
test loss: 0.19370371103286743, acc: 0.9464
train loss: 0.1122715100646019, acc: 0.9692
--------------------------------------------------------------------------------
test loss: 0.16738709807395935, acc: 0.9538
train loss: 0.08012499660253525, acc: 0.9787
--------------------------------------------------------------------------------
test loss: 0.16035553812980652, acc: 0.9575
train loss: 0.06216369569301605, acc: 0.9838
--------------------------------------------------------------------------------
test loss: 0.15690605342388153, acc: 0.9587
train loss: 0.04842701926827431, acc: 0.9877
--------------------------------------------------------------------------------
test loss: 0.1597040444612503, acc: 0.9586
train loss: 0.03863723576068878, acc: 0.9909
--------------------------------------------------------------------------------
test loss: 0.16320295631885529, acc: 0.9593
train loss: 0.031261660158634186, acc: 0.9933
--------------------------------------------------------------------------------
test loss: 0.1675170212984085, acc: 0.959
train loss: 0.02533782459795475, acc: 0.9948
--------------------------------------------------------------------------------
test loss: 0.17022284865379333, acc: 0.9592
train loss: 0.020637042820453644, acc: 0.9962
--------------------------------------------------------------------------------
'''

 rnn中pad和pack的使用

torch.nn.utils.rnn.pad_sequence()
torch.nn.utils.rnn.pack_padded_sequence()
torch.nn.utils.rnn.pad_packed_sequence()

　　LSTM (BiLSTM) 分词

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import os
import time
import matplotlib.pyplot as plt

np.random.seed(1)
torch.manual_seed(1)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
print('device:', device)
device = torch.device(device)

small = 50000
training = True
show_loss = True
N_epoch = 10
Batch_size = 256
show_epoch = 1
H = 128
em_dim = 100
lr = 1e-2
fold=0
net = 'BiLSTM'


words = {}
max_seq = 0

label2idx = {'B': 0, 'I': 1, 'S': 2}
idx2label = {0: 'B', 1: 'I', 2: 'S'}

Sentences = []
Sentences_label = []
Sentences_origin = []

with open('zhihu.txt', mode='r', encoding='utf8') as f:
    lines = f.readlines()
    print('len(lines):', len(lines))
    for idx, line in enumerate(lines):
        # print(line)
        line = line.split()
        tmp = []
        mmp = []
        for word in line:
            if len(word)==1:
                mmp.append(label2idx['S'])
            else:
                mmp.append(label2idx['B'])
                for _ in range(1, len(word)):
                    mmp.append(label2idx['I'])

            for w in word:
                if w in words:
                    words[w] += 1
                else:
                    words[w] = 1
                tmp.append(w)
        Sentences.append(tmp)
        Sentences_label.append(mmp)
        max_seq = max(max_seq, len(tmp))
        assert len(mmp)==len(tmp)
        # print(tmp)
        if idx > small:
            break

Sentences.sort(key=lambda x: len(x), reverse=True)
Sentences_label.sort(key=lambda x: len(x), reverse=True)
# for sentence in Sentences:
#     print(sentence)
print('len(words):', len(words))
# print(words)
print('max_seq len:', max_seq)

# print(words)
word2index = {word: idx + 1 for idx, word in enumerate(words.keys())}
indx2word = {idx + 1: word for idx, word in enumerate(words.keys())}


voc_size = len(words) + 1
word2index['<pad>'] = 0
indx2word[0] = '<pad>'

Sentences_idx, Sentences_len = [], []
for sentence in Sentences:
    tmp=[]
    for w in sentence:
        tmp.append(word2index[w])
    Sentences_idx.append(torch.LongTensor(tmp))
    Sentences_len.append(len(tmp))
    # print(tmp)
Sentences_idx = torch.nn.utils.rnn.pad_sequence(Sentences_idx,batch_first=True)
# print('-' * 80)
# print(Sentences_idx.size())
# print(Sentences_idx)


# print('-' * 80)
Sentences_label_idx = []
for i, sentences_label in enumerate(Sentences_label):
    tmp = torch.LongTensor(sentences_label)
    Sentences_label_idx.append(tmp)
    # print(Sentences[i])
    # # print(lines[i])
    # print(tmp)
    assert len(tmp) == len(Sentences[i])
Sentences_label_idx = torch.nn.utils.rnn.pad_sequence(Sentences_label_idx,batch_first=True,padding_value=0)
# print('Sentences_label_idx:')
# print(Sentences_label_idx)

# a = torch.tensor(1.0)
# print(a)
# print(a.size())
class MyDataSet(Dataset):
    def __init__(self, data, lens, labels):
        self.data = data
        self.lens = lens
        self.labels = labels
    def __getitem__(self, idx):
        now_data = self.data[idx]
        now_len = self.lens[idx]
        now_mask = []
        now_label = self.labels[idx]
        for i in range(len(now_data)):
            t = 1.0 if i < now_len else 0.0
            now_mask.append(t)
        now_mask = torch.Tensor(now_mask)
        return now_data, now_len, now_mask, now_label
    def __len__(self):
        return len(self.data)

class FenCi_Zqx(nn.Module):
    def __init__(self, voc_size, em_dim, H):
        super(FenCi_Zqx, self).__init__()
        self.emd = nn.Embedding(num_embeddings=voc_size,embedding_dim=em_dim)
        if net == 'LSTM':
            self.rnn = nn.LSTM(input_size=em_dim,hidden_size=H,num_layers=1,batch_first=True)
            self.linear = nn.Linear(in_features=H,out_features=3)
        if net == 'BiLSTM':
            self.rnn = nn.LSTM(input_size=em_dim, hidden_size=H, num_layers=1, batch_first=True,bidirectional=True)
            self.linear = nn.Linear(in_features=2*H, out_features=3)
    def forward(self, sentence, sentence_len=None, mask=None):
        emd = self.emd(sentence) #  (batch, seq_len, em_dim)
        all_h, (h, c) = self.rnn(emd) # LSTM: (batch, seq_len, H)   BiLSTM: (batch, seq_len, 2*H)
        # print('emd size:', emd.size())
        # print('all_h.size():', all_h.size())
        # out = all_h.view(-1, all_h.size(2))  # (batch * seq_len, H)
        out = self.linear(all_h).view(emd.size(0), emd.size(1), 3) # (batch, seq_len, 3)
        # print('out size:', out.size())
        return out

Sentences_len = torch.Tensor(Sentences_len)
train_idx = [i for i in range(len(Sentences_len)) if i % 5 == fold]
test_idx = [i for i in range(len(Sentences_len)) if i % 5 != fold]
print(train_idx, '
', test_idx)
Train_Sentences_idx, Train_Sentences_len, Train_Sentences_label_idx = 
    Sentences_idx[train_idx], Sentences_len[train_idx], Sentences_label_idx[train_idx]
Test_Sentences_idx, Test_Sentences_len, Test_Sentences_label_idx = 
    Sentences_idx[test_idx], Sentences_len[test_idx], Sentences_label_idx[test_idx]
Train_data = MyDataSet(Train_Sentences_idx, Train_Sentences_len, Train_Sentences_label_idx)
Train_data_loader = DataLoader(dataset=Train_data, batch_size=Batch_size, shuffle=True)
Test_data = MyDataSet(Test_Sentences_idx, Test_Sentences_len, Test_Sentences_label_idx)
Test_data_loader = DataLoader(dataset=Test_data, batch_size=Batch_size, shuffle=False)

model = FenCi_Zqx(voc_size=voc_size, em_dim=em_dim, H=H)
loss_fn = nn.CrossEntropyLoss(reduction='none')
opt = torch.optim.Adam(model.parameters(), lr=lr)


print('Sentences_idx, Sentences_len, Sentences_label_idx shape')
print(len(Sentences_idx), len(Sentences_len), len(Sentences_label_idx))
print(Sentences_idx.size(), Sentences_len.size(), Sentences_label_idx.size())
print(Sentences_idx.shape, Sentences_len.shape, Sentences_label_idx.shape)
print('#' * 60)
print(model)

def valid(model):
    # model.to(device)
    # model.eval()
    with torch.no_grad():
        avg_loss = 0
        cnt=0
        for batch_step, (now_data, now_len, now_mask, now_label) in enumerate(Test_data_loader):
            cnt += 1
            now_data, now_len, now_mask, now_label = 
                now_data.to(device), now_len.to(device), now_mask.to(device), now_label.to(device)
            out = model(now_data, now_len, now_mask)
            out = out.view(-1, 3)
            now_mask = now_mask.view(-1)
            now_label = now_label.view(-1)
            loss = loss_fn(out, now_label)
            # print('loss size:', loss.size())
            # print(out.size(), now_label.size(), now_mask.size())
            loss = torch.mean(loss * now_mask)
            avg_loss += loss.item()
            # print('loss size:', loss.size())
            # print('loss:', loss.item())

        avg_loss /= cnt
        return avg_loss


def train(model):
    print('start training:')
    model.to(device)
    time_st_global = time.time()
    Train_loss,Valid_loss = [], []
    for epoch in range(N_epoch):
        time_st_epoch = time.time()
        avg_loss = 0
        cnt = 0
        for batch_step, (now_data, now_len, now_mask, now_label) in enumerate(Train_data_loader):
            cnt += 1
            now_data, now_len, now_mask, now_label = 
                now_data.to(device), now_len.to(device), now_mask.to(device), now_label.to(device)
            out = model(now_data, now_len, now_mask)
            out = out.view(-1, 3)
            now_mask = now_mask.view(-1)
            now_label = now_label.view(-1)
            loss = loss_fn(out, now_label)
            # print('loss size:', loss.size())
            # print(out.size(), now_label.size(), now_mask.size())
            loss = torch.mean(loss * now_mask)
            avg_loss += loss.item()
            # print('loss size:', loss.size())
            # print('loss:', loss.item())

            opt.zero_grad()
            loss.backward()
            opt.step()
        avg_loss /= cnt
        valid_avg_loss = valid(model)
        print('#' * 80)
        print('epoch:{}, steps: {}, train avg loss: {} -- valid avg loss : {} '.format(epoch, cnt, avg_loss, valid_avg_loss))
        print('epoch time {:.2f} s, total time {:.2f} s'.format(time.time() - time_st_epoch,
                                                                time.time() - time_st_global))
        if len(Valid_loss)==0 or valid_avg_loss < min(Valid_loss):
            if not os.path.exists(check_path):
                os.makedirs(check_path)
            torch.save(model.state_dict(), filepath)

        Train_loss.append(avg_loss)
        Valid_loss.append(valid_avg_loss)

    if show_loss:
        plt.figure()
        plt.plot(Train_loss,label='Train loss')
        plt.plot(Valid_loss, label='Valid loss')
        plt.legend()
        plt.savefig('Train_Valid_loss' + net + '.png')
        # plt.show()
    return model

        # break

check_path = './Checkpoints/'
filepath = check_path + 'p03_Fenci_state_dict_' + net + ' .pkl'

if training:
    model = train(model)


# 模型恢复
model.load_state_dict(torch.load(filepath))


test_words = [
    '我是中国人，我爱祖国',
    '独行侠队的球员们承诺每天为达拉斯地区奋战在抗疫一线的工作人员们提供餐食',
    '汤普森太爱打球,不能出场让他很煎熬',
    '这个赛季对克莱来说非常艰难，他太热爱打篮球了，无法上场让他很受打击。',
    '克莱和斯蒂芬会处在极佳的状态，准备好比赛。',
    '勇士已经证明了他们也是一支历史级别的球队，维金斯在稍强于巴恩斯的前提下，仍然算得上是三号位上一位合格的替代者'
]

np.set_printoptions(precision=3, suppress=True)
model.cpu()
model.eval()
for word in test_words:
    print('-' * 80)
    print('test word : {}'.format(word))
    word_idx = [word2index[w] for w in word]
    word_idx = torch.LongTensor([word_idx])
    # print('word_idx.size():', word_idx.size())
    # word_idx.to(device)
    out = model(word_idx)
    # print('out.size():', out.size())
    out = out.squeeze(0).cpu().detach().numpy()
    # print('out.shape():', out.shape)
    # print(out)
    out_label = np.argmax(out, axis=1)
    # print(out_label)

    for i, w in enumerate(word):
        print('{} -> {} -> {}'.format(w, idx2label[out_label[i]], out_label[i]))

print('end!!!')
'''
test word : 勇士已经证明了他们也是一支历史级别的球队，维金斯在稍强于巴恩斯的前提下，仍然算得上是三号位上一位合格的替代者
勇 -> B -> 0
士 -> I -> 1
已 -> B -> 0
经 -> I -> 1
证 -> B -> 0
明 -> I -> 1
了 -> S -> 2
他 -> S -> 2
们 -> I -> 1
也 -> S -> 2
是 -> S -> 2
一 -> S -> 2
支 -> S -> 2
历 -> B -> 0
史 -> I -> 1
级 -> B -> 0
别 -> I -> 1
的 -> S -> 2
球 -> B -> 0
队 -> I -> 1
， -> S -> 2
维 -> B -> 0
金 -> I -> 1
斯 -> I -> 1
在 -> S -> 2
稍 -> B -> 0
强 -> B -> 0
于 -> I -> 1
巴 -> B -> 0
恩 -> I -> 1
斯 -> I -> 1
的 -> S -> 2
前 -> B -> 0
提 -> I -> 1
下 -> S -> 2
， -> S -> 2
仍 -> B -> 0
然 -> I -> 1
算 -> I -> 1
得 -> I -> 1
上 -> S -> 2
是 -> S -> 2
三 -> S -> 2
号 -> S -> 2
位 -> S -> 2
上 -> B -> 0
一 -> B -> 0
位 -> S -> 2
合 -> B -> 0
格 -> I -> 1
的 -> S -> 2
替 -> B -> 0
代 -> I -> 1
者 -> I -> 1

'''

 　　

　　HMM分词

import numpy as np
import os
import time
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
np.random.seed(1)

small = 50000
training = True
show_loss = True
N_epoch = 10
Batch_size = 256
show_epoch = 1
H = 128
em_dim = 100
lr = 1e-2
fold=5
net = 'HMM'
check_path = './Checkpoints/'
filepath = check_path + 'p03_Fenci_state_dict_' + net + ' .pkl'

words = {}
max_seq = 0

# label2idx = {'B': 0, 'I': 1, 'S': 2, 'BOS': 3, 'EOS': 4}
# idx2label = {0: 'B', 1: 'I', 2: 'S', 3: 'BOS', 4: 'EOS'}

Sentences = []
Sentences_tag = []

with open('zhihu.txt', mode='r', encoding='utf8') as f:
    lines = f.readlines()
    print('len(lines):', len(lines))
    for idx, line in enumerate(lines):
        # print(line)
        line = line.split()
        tmp = []
        mmp = []
        for word in line:
            if len(word)==1:
                mmp.append('S')
            else:
                mmp.append('B')
                for _ in range(1, len(word)):
                    mmp.append('I')
            for w in word:
                if w in words:
                    words[w] += 1
                else:
                    words[w] = 1
                tmp.append(w)
        Sentences.append(tmp)   # 存下以字单位的sentence
        Sentences_tag.append(mmp)  #存下每个sentence中每个字对应的BIS标签
        max_seq = max(max_seq, len(tmp))
        assert len(mmp)==len(tmp)
        assert len(tmp)> 0 # 判断是否存在空的sentence
        # print(tmp)
        if idx > small:
            break

print('len(words):', len(words))
print('max_seq len:', max_seq)

for idx, sentence in enumerate(Sentences):
    print('-' * 80)
    print(sentence)
    print(Sentences_tag[idx])
    if idx > 5:
        break

Train_Sentences, Train_Sentences_tag = [], []
Valid_Sentences, Valid_Sentences_tag = [], []
for i in range(len(Sentences)):
    if i % fold:
        Train_Sentences.append(Sentences[i])
        Train_Sentences_tag.append(Sentences_tag[i])
    else:
        Valid_Sentences.append(Sentences[i])
        Valid_Sentences_tag.append(Sentences_tag[i])

loss_fn = nn.CrossEntropyLoss(reduction='none')

def train(Train_Sentences, Train_Sentences_tag):
    N = len(Train_Sentences)
    tag, tag2word, tag2tag = {}, {}, {}
    tag['BOS'] = N
    tag['EOS'] = N
    for i in range(N):
        sentence = Train_Sentences[i]
        sentence_tag = Train_Sentences_tag[i]
        n = len(sentence)
        assert len(sentence) == len(sentence_tag)
        assert n > 0
        if ('BOS', sentence_tag[0]) in tag2tag:
            tag2tag[('BOS', sentence_tag[0])] += 1
        else:
            tag2tag[('BOS', sentence_tag[0])] = 1
        if (sentence_tag[-1], 'EOS') in tag2tag:
            tag2tag[(sentence_tag[-1], 'EOS')] += 1
        else:
            tag2tag[(sentence_tag[-1], 'EOS')] = 1

        for i in range(n):
            tg, w = sentence_tag[i], sentence[i]
            if tg in tag:
                tag[tg] += 1
            else:
                tag[tg] = 1
            if (tg, w) in tag2word:
                tag2word[(tg, w)] += 1
            else:
                tag2word[(tg, w)] = 1

            if i < n - 1:
                next_tg = sentence_tag[i + 1]
                if (tg, next_tg) in tag2tag:
                    tag2tag[(tg, next_tg)] += 1
                else:
                    tag2tag[(tg, next_tg)] = 1
    Prob_tag2tag, Prob_tag2word = {}, {}
    for tg1, tg2 in tag2tag.keys():
        Prob_tag2tag[(tg1, tg2)] = 0.0 + tag2tag[(tg1, tg2)] / tag[tg1]
    for tg, w in tag2word.keys():
        Prob_tag2word[(tg, w)] = 0.0 + tag2word[(tg, w)] / tag[tg]
    # print('tag:{} 
tag2word:{} 
tag2tag:{} 
'.format(tag, tag2word, tag2tag))
    print('tag:{} 
tag2word:{} 
tag2tag:{} 
'.format(len(tag), len(tag2word), len(tag2tag)))
    # print('
Prob_tag2word:{} 
Prob_tag2tag:{} 
'.format(Prob_tag2word, Prob_tag2tag))
    print('
Prob_tag2word:{} 
Prob_tag2tag:{} 
'.format(len(Prob_tag2word), len(Prob_tag2tag)))
    return tag, tag2word, tag2tag, Prob_tag2tag, Prob_tag2word
Tag, Tag2word, Tag2tag, Prob_tag2tag, Prob_tag2word = train(Train_Sentences, Train_Sentences_tag)

def predict_tag(sentence, True_sentence_tag=None):
    n = len(sentence)
    tags = ['B', 'I', 'S', 'BOS', 'EOS']
    dp = [{'B': 0.0, 'I': 0.0, 'S': 0.0, 'BOS': 0.0, 'EOS': 0.0} for _ in range(n + 1)]
    pre_tag = [{'B': None, 'I': None, 'S': None, 'BOS': None, 'EOS': None} for _ in range(n + 1)]
    for t in range(n):
        w = sentence[t]
        # print('w:', w)
        for tg in tags:
            prob_tag2word = 1e-9 if (tg, w) not in Prob_tag2word else Prob_tag2word[(tg, w)]
            if t == 0:
                prob_tag2tag = 1e-9 if ('BOS', tg) not in Prob_tag2tag else Prob_tag2tag[('BOS', tg)]
                dp[t][tg] = np.log(prob_tag2tag) + np.log(prob_tag2word)
                pre_tag[t][tg] = 'BOS'
            else:
                max_prob = None
                best_pre_tag = None
                for pre_tg in tags:
                    prob_tag2tag = 1e-9 if (pre_tg, tg) not in Prob_tag2tag else Prob_tag2tag[(pre_tg, tg)]
                    tmp = dp[t - 1][pre_tg] + np.log(prob_tag2tag) + np.log(prob_tag2word)
                    if max_prob == None or max_prob < tmp:
                        max_prob = tmp
                        best_pre_tag = pre_tg
                dp[t][tg] = max_prob
                pre_tag[t][tg] = best_pre_tag

    max_prob = None
    best_pre_tag = None
    tg = 'EOS'
    for pre_tg in tags:
        prob_tag2tag = 1e-9 if (pre_tg, tg) not in Prob_tag2tag else Prob_tag2tag[(pre_tg, tg)]
        tmp = dp[n - 1][pre_tg] + np.log(prob_tag2tag)
        if max_prob == None or max_prob < tmp:
            max_prob = tmp
            best_pre_tag = pre_tg
    dp[n][tg] = max_prob
    pre_tag[n][tg] = best_pre_tag

    ans_tag = []
    t = n

    # print('#' * 80)
    # print('sentence:', sentence)
    # print('True sentence tag:', True_sentence_tag)
    # print('len(sentence):', len(sentence))
    # print('n:', n)
    if True_sentence_tag is not None:
        True_sentence_tag.append('EOS')
    sss = sentence + ['END']
    while pre_tag[t][tg] is not None:
        if True_sentence_tag is None:
            # print('t: {}, pre_tag[t][tg]: {} -> tg: {} -- word:{}'.format(
            #     t,  pre_tag[t][tg], tg, sss[t]))
            pass
        else:
            assert len(True_sentence_tag) == n + 1, (n, len(True_sentence_tag))
            print('t: {}, pre_tag[t][tg]: {} -> tg: {}  -- True tag: {}, -- word: {}'.format(
                t, pre_tag[t][tg], tg, True_sentence_tag[t], sss[t]))

        ans_tag = [pre_tag[t][tg]] + ans_tag
        tg = pre_tag[t][tg]
        t = t - 1

    return ans_tag[1:]  # 去掉BOS

# predict_tag(sentence=Sentences[0], True_sentence_tag=Sentences_tag[0])
predict_tag(sentence=Sentences[0], True_sentence_tag=None)


def fenci_example():

    test_sentences = [
        '我是中国人，我爱祖国',
        '独行侠队的球员们承诺每天为达拉斯地区奋战在抗疫一线的工作人员们提供餐食',
        '汤普森太爱打球,不能出场让他很煎熬',
        '这个赛季对克莱来说非常艰难，他太热爱打篮球了，无法上场让他很受打击。',
        '克莱和斯蒂芬会处在极佳的状态，准备好比赛。',
        '勇士已经证明了他们也是一支历史级别的球队，维金斯在稍强于巴恩斯的前提下，仍然算得上是三号位上一位合格的替代者'
    ]

    np.set_printoptions(precision=3, suppress=True)

    for sentence in test_sentences:
        print('-' * 80)
        print('test word : {}'.format(sentence))
        sentence = [w for w in sentence]
        sentence_tag = predict_tag(sentence)
        # predict_tag(sentence=Sentences[0], True_sentence_tag=None)
        for i, w in enumerate(sentence):
            print('{} -> {}'.format(w, sentence_tag[i]))
fenci_example()

print('end!!!')
'''
test word : 克莱和斯蒂芬会处在极佳的状态，准备好比赛。
克 -> B
莱 -> I
和 -> S
斯 -> B
蒂 -> I
芬 -> I
会 -> S
处 -> B
在 -> I
极 -> B
佳 -> I
的 -> S
状 -> B
态 -> I
， -> S
准 -> B
备 -> I
好 -> S
比 -> B
赛 -> I
。 -> S
'''

　　CRF分词， ps:只训练了一个epoch，不知道为什么，中间梯度爆炸了， 后来查了下，应该要用logsumexp函数，见我写的 CRF layer

import numpy as np
import os
import time
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
np.random.seed(1)

small = 50000
training = True
show_loss = True
show_acc = True
reused_W = False
N_epoch = 1
Batch_size = 256
show_epoch = 10
H = 128
em_dim = 100
lr = 1e-2
fold=5
net = 'HMM'
check_path = './Checkpoints/'
filepath = check_path + 'W_crf.npy'
regulization = 0

words = {}
max_seq = 0

# label2idx = {'B': 0, 'I': 1, 'S': 2, 'BOS': 3, 'EOS': 4}
# idx2label = {0: 'B', 1: 'I', 2: 'S', 3: 'BOS', 4: 'EOS'}
tag_num = 5
Sentences = []
Sentences_tag = []

with open('zhihu.txt', mode='r', encoding='utf8') as f:
    lines = f.readlines()
    print('len(lines):', len(lines))
    for idx, line in enumerate(lines):
        # print(line)
        line = line.split()
        tmp = []
        mmp = []
        for word in line:
            if len(word)==1:
                mmp.append('S')
            else:
                mmp.append('B')
                for _ in range(1, len(word)):
                    mmp.append('I')
            for w in word:
                if w in words:
                    words[w] += 1
                else:
                    words[w] = 1
                tmp.append(w)
        Sentences.append(tmp)   # 存下以字单位的sentence
        Sentences_tag.append(mmp)  #存下每个sentence中每个字对应的BIS标签
        max_seq = max(max_seq, len(tmp))
        assert len(mmp)==len(tmp)
        assert len(tmp)> 0 # 判断是否存在空的sentence
        # print(tmp)
        if idx > small:
            break

print('len(words):', len(words))
print('max_seq len:', max_seq)

for idx, sentence in enumerate(Sentences):
    print('-' * 80)
    print(sentence)
    print(Sentences_tag[idx])
    if idx > 5:
        break

Train_Sentences, Train_Sentences_tag = [], []
Valid_Sentences, Valid_Sentences_tag = [], []
for i in range(len(Sentences)):
    if i % fold:
        Train_Sentences.append(Sentences[i])
        Train_Sentences_tag.append(Sentences_tag[i])
    else:
        Valid_Sentences.append(Sentences[i])
        Valid_Sentences_tag.append(Sentences_tag[i])

loss_fn = nn.CrossEntropyLoss(reduction='none')

def predict_tag(W, feat_pair2idx, sentence, True_sentence_tag=None):
    n = len(sentence)
    tags = ['B', 'I', 'S', 'BOS', 'EOS']
    dp = [{'B': 0.0, 'I': 0.0, 'S': 0.0, 'BOS': 0.0, 'EOS': 0.0} for _ in range(n + 1)]
    pre_tag = [{'B': None, 'I': None, 'S': None, 'BOS': None, 'EOS': None} for _ in range(n + 1)]
    for t in range(n):
        w = sentence[t]
        # print('w:', w)
        for tg in tags:
            feat_tag2word = -1e9 if (tg, w) not in feat_pair2idx else W[feat_pair2idx[(tg, w)]]
            if t == 0:
                feat_tag2tag = -1e9 if ('BOS', tg) not in feat_pair2idx else W[feat_pair2idx[('BOS', tg)]]
                dp[t][tg] = feat_tag2word + feat_tag2tag
                pre_tag[t][tg] = 'BOS'
            else:
                max_prob = None
                best_pre_tag = None
                for pre_tg in tags:
                    feat_tag2tag = -1e9 if (pre_tg, tg) not in feat_pair2idx else W[feat_pair2idx[(pre_tg, tg)]]
                    tmp = dp[t - 1][pre_tg] + feat_tag2tag + feat_tag2word
                    if max_prob == None or max_prob < tmp:
                        max_prob = tmp
                        best_pre_tag = pre_tg
                dp[t][tg] = max_prob
                pre_tag[t][tg] = best_pre_tag

    max_prob = None
    best_pre_tag = None
    tg = 'EOS'
    for pre_tg in tags:
        feat_tag2tag = -1e9 if (pre_tg, tg) not in feat_pair2idx else W[feat_pair2idx[(pre_tg, tg)]]
        tmp = dp[n - 1][pre_tg] + feat_tag2tag
        if max_prob == None or max_prob < tmp:
            max_prob = tmp
            best_pre_tag = pre_tg
    dp[n][tg] = max_prob
    pre_tag[n][tg] = best_pre_tag

    ans_tag = []
    t = n

    # print('#' * 80)
    # print('sentence:', sentence)
    # print('True sentence tag:', True_sentence_tag)
    # print('len(sentence):', len(sentence))
    # print('n:', n)

    if True_sentence_tag is not None:
        True_sentence_tag.append('EOS')
    sss = sentence + ['END']
    while pre_tag[t][tg] is not None:
        if True_sentence_tag is None:
            # print('t: {}, pre_tag[t][tg]: {} -> tg: {} -- word:{}'.format(
            #     t,  pre_tag[t][tg], tg, sss[t]))
            pass
        else:
            assert len(True_sentence_tag) == n + 1, (n, len(True_sentence_tag))
            print('t: {}, pre_tag[t][tg]: {} -> tg: {}  -- True tag: {}, -- word: {}'.format(
                t, pre_tag[t][tg], tg, True_sentence_tag[t], sss[t]))

        ans_tag = [pre_tag[t][tg]] + ans_tag
        tg = pre_tag[t][tg]
        t = t - 1

    return ans_tag[1:]  # 去掉BOS


def cal_grad_w(W, feat_pair2idx, feat_num, xn, yn):
    """
    O(W, xn, yn) = log p(yn|xn) = log exp(W Phi(xn, yn)) / Sigma exp(W Phi(xn, y'))
        =W Phi(xn, yn) - log Sigma exp(W Phi(xn, y'))
        = W Phi(xn, yn) - log Z(xn)
    W_grad = Phi(xn, yn) - 1 / Z(xn)  * Sigma exp(W Phi(xn, y')) Phi(xn, y')
    然后利用viterbi算法进行求解，实际上就是 O( 序列长度 * tag种类数 ** 2)的动态规划算法
    我们用到两个东西：
        1. Z_i(t)：表示t时刻为止，tag是i的所有路径的概率之和，
                i.e.,  Sigma exp(W Phi(xn(1:t), y'(1:t))) 且y(t) = tag i
        2. P_i(t): 表示t时刻为止，tag是i的所有路径的【加权】（Phi(xn(1:t), y'(1:t))）概率之和，
                i.e.,  Sigma exp(W Phi(xn(1:t), y'(1:t))) Phi(xn(1:t), y'(1:t))
    具体状态转移方程见代码， 关键是
        P_i(t + 1) = exp(W Phi(xn(1:t+1), y'(1:t+1))) Phi(xn(1:t+1), y'(1:t+1))
            = exp(W Phi_t) exp (W delta_Phi) * (Phi_t +delta_Phi)
            =  Sigma_{y'(t)}  (exp(W Phi_t)Phi_t + delta_Phi) * exp (W delta_Phi)
        Z_i(t + 1) = Sigma_{y'(t)} Z_i(t) * exp (W delta_Phi)
    为了数值稳定，可以用log_P和log_Z进行更新
    如果看不懂上面，可以参考下面的链接（可能还是比较模糊），最好自己推导一边
    链接1：https://blog.csdn.net/qq_42189083/article/details/89350890
    链接2：https://blog.csdn.net/weixin_30014549/article/details/52850638
    """
    tags = ['B', 'I', 'S', 'BOS', 'EOS']
    Phi = np.zeros(feat_num)
    pre_P = np.zeros(shape=[5, feat_num])
    pre_Z = np.zeros(shape=[5,])
    n = len(xn)

    pre_tag = 'BOS'
    for i in range(n):
        word, tag = xn[i], yn[i]
        tag2tag_id = feat_pair2idx[(pre_tag, tag)]
        tag2word_id = feat_pair2idx[(tag, word)]
        Phi[tag2tag_id] += 1
        Phi[tag2word_id] += 1
        pre_tag = tag

    for i in range(n):
        word = xn[i]

        P = np.zeros(shape=[5, feat_num])
        Z = np.zeros(shape=[5, ])
        flag = 0
        for j, tag in enumerate(tags):
            for k, pre_tag in enumerate(tags):
                if i==0 and pre_tag != 'BOS':
                    continue
                deta_phi = np.zeros(feat_num)
                tag2tag = (pre_tag, tag)
                tag2word = (tag, word)
                if tag2tag not in feat_pair2idx:
                    continue
                if tag2word not in feat_pair2idx:
                    continue
                flag = 1
                tag2tag_id = feat_pair2idx[tag2tag]
                tag2word_id = feat_pair2idx[tag2word]
                deta_phi[tag2tag_id] += 1
                deta_phi[tag2word_id] += 1

                # exp_w_delta_phi = np.exp(np.sum(W * deta_phi))
                exp_w_delta_phi = np.exp(W[tag2tag_id] + W[tag2word_id])

                if i == 0 and pre_tag == 'BOS':
                    pre_Z[k] = 1
                P[j] += (pre_P[k] + pre_Z[k] * deta_phi) * exp_w_delta_phi
                Z[j] += pre_Z[k] * exp_w_delta_phi

                # print('P[j, tag2tag_id]:{}, P[j, tag2word_id]:{}'.format(P[j, tag2tag_id], P[j, tag2word_id]))
        pre_P = P.copy()
        pre_Z = Z.copy()
        # print('word: {}, flag: {}'.format(word, flag))

    P = np.zeros(shape=[feat_num, ])
    Z = 0.0
    tag = 'EOS'
    for k, pre_tag in enumerate(tags):
        deta_phi = np.zeros(feat_num)
        tag2tag = (pre_tag, tag)
        if tag2tag not in feat_pair2idx:
            continue
        tag2tag_id = feat_pair2idx[tag2tag]
        deta_phi[tag2tag_id] += 1
        # exp_w_delta_phi = np.exp(np.sum(W * deta_phi))
        exp_w_delta_phi = np.exp(W[tag2tag_id])

        P += (pre_P[k] + pre_Z[k] * deta_phi) * exp_w_delta_phi
        Z += pre_Z[k] * exp_w_delta_phi
    # print('pre_P: {}
pre_Z: {}
'.format(pre_P, pre_Z))
    # print('sum(Phi): {}
P:{}
Z:{}'.format(np.sum(Phi), P, Z))
    # print('WPhi: {}, exp(WPhi):{}'.format(np.sum(W * Phi), np.exp(np.sum(W * Phi))))
    # print('Phi - P / Z:', Phi - P / Z)
    W_grad = Phi - P / Z
    return - W_grad + regulization * W

def cal_grad_w_log_version(W, feat_pair2idx, feat_num, xn, yn):
    """
    O(W, xn, yn) = log p(yn|xn) = log exp(W Phi(xn, yn)) / Sigma exp(W Phi(xn, y'))
        =W Phi(xn, yn) - log Sigma exp(W Phi(xn, y'))
        = W Phi(xn, yn) - log Z(xn)
    W_grad = Phi(xn, yn) - 1 / Z(xn)  * Sigma exp(W Phi(xn, y')) Phi(xn, y')
    然后利用viterbi算法进行求解，实际上就是 O( 序列长度 * tag种类数 ** 2)的动态规划算法
    我们用到两个东西：
        1. Z_i(t)：表示t时刻为止，tag是i的所有路径的概率之和，
                i.e.,  Sigma exp(W Phi(xn(1:t), y'(1:t))) 且y(t) = tag i
        2. P_i(t): 表示t时刻为止，tag是i的所有路径的【加权】（Phi(xn(1:t), y'(1:t))）概率之和，
                i.e.,  Sigma exp(W Phi(xn(1:t), y'(1:t))) Phi(xn(1:t), y'(1:t))
    具体状态转移方程见代码， 关键是
        P_i(t + 1) = exp(W Phi(xn(1:t+1), y'(1:t+1))) Phi(xn(1:t+1), y'(1:t+1))
            = exp(W Phi_t) exp (W delta_Phi) * (Phi_t +delta_Phi)
            =  Sigma_{y'(t)}  (exp(W Phi_t)Phi_t + delta_Phi) * exp (W delta_Phi)
        Z_i(t + 1) = Sigma_{y'(t)} Z_i(t) * exp (W delta_Phi)
    为了数值稳定，可以用log_P和log_Z进行更新
    如果看不懂上面，可以参考下面的链接（可能还是比较模糊），最好自己推导一边
    链接1：https://blog.csdn.net/qq_42189083/article/details/89350890
    链接2：https://blog.csdn.net/weixin_30014549/article/details/52850638
    """
    tags = ['B', 'I', 'S', 'BOS', 'EOS']
    Phi = np.zeros(feat_num)
    log_pre_P = np.zeros(shape=[5, feat_num])
    log_pre_Z = np.zeros(shape=[5,])
    n = len(xn)

    pre_tag = 'BOS'
    for i in range(n):
        word, tag = xn[i], yn[i]
        tag2tag_id = feat_pair2idx[(pre_tag, tag)]
        tag2word_id = feat_pair2idx[(tag, word)]
        Phi[tag2tag_id] += 1
        Phi[tag2word_id] += 1
        pre_tag = tag

    for i in range(n):
        word = xn[i]

        log_P = np.zeros(shape=[5, feat_num]) + 1e-9
        log_Z = np.zeros(shape=[5, ]) + 1e-9
        flag = 0
        for j, tag in enumerate(tags):
            for k, pre_tag in enumerate(tags):
                if i==0 and pre_tag != 'BOS':
                    continue
                deta_phi = np.zeros(feat_num)
                tag2tag = (pre_tag, tag)
                tag2word = (tag, word)
                if tag2tag not in feat_pair2idx:
                    continue
                if tag2word not in feat_pair2idx:
                    continue
                flag = 1
                tag2tag_id = feat_pair2idx[tag2tag]
                tag2word_id = feat_pair2idx[tag2word]
                deta_phi[tag2tag_id] += 1
                deta_phi[tag2word_id] += 1

                # exp_w_delta_phi = np.exp(np.sum(W * deta_phi))
                exp_w_delta_phi = np.exp(W[tag2tag_id] + W[tag2word_id])

                if i == 0 and pre_tag == 'BOS':
                    log_pre_Z[k] = 0
                log_P[j] += (np.exp(log_pre_P[k]) + np.exp(log_pre_Z[k]) * deta_phi) * exp_w_delta_phi
                log_Z[j] += np.exp(log_pre_Z[k]) * exp_w_delta_phi

                # print('P[j, tag2tag_id]:{}, P[j, tag2word_id]:{}'.format(log_P[j, tag2tag_id], log_P[j, tag2word_id]))
        log_P = np.log(log_P)
        log_Z = np.log(log_Z)
        log_pre_P = log_P.copy()
        log_pre_Z = log_Z.copy()
        # print('word: {}, flag: {}'.format(word, flag))

    log_P = np.zeros(shape=[feat_num, ])
    log_Z = 0.0
    tag = 'EOS'
    for k, pre_tag in enumerate(tags):
        deta_phi = np.zeros(feat_num)
        tag2tag = (pre_tag, tag)
        if tag2tag not in feat_pair2idx:
            continue
        tag2tag_id = feat_pair2idx[tag2tag]
        deta_phi[tag2tag_id] += 1
        # exp_w_delta_phi = np.exp(np.sum(W * deta_phi))
        exp_w_delta_phi = np.exp(W[tag2tag_id])

        log_P += (np.exp(log_pre_P[k]) + np.exp(log_pre_Z[k]) * deta_phi) * exp_w_delta_phi
        log_Z += np.exp(log_pre_Z[k]) * exp_w_delta_phi
    # print('pre_P: {}
pre_Z: {}
'.format(pre_P, pre_Z))
    # print('sum(Phi): {}
P:{}
Z:{}'.format(np.sum(Phi), P, Z))
    # print('WPhi: {}, exp(WPhi):{}'.format(np.sum(W * Phi), np.exp(np.sum(W * Phi))))
    # print('Phi - P / Z:', Phi - P / Z)
    W_grad = Phi - log_P / log_Z
    return - W_grad + regulization * W

def evaluate(W, feat_pair2idx, Sentences_, Sentences_tag_):
    cnt_correct_tag, cnt_total_tag = 0.0, 0.0
    for i, sentence in enumerate(Sentences_):
        sentence_tag = Sentences_tag_[i]
        sentence_tag_pred = predict_tag(W, feat_pair2idx, sentence)
        assert len(sentence_tag) == len(sentence_tag_pred)
        # predict_tag(sentence=Sentences[0], True_sentence_tag=None)
        # print('sentence_tag == sentence_tag_pred:', [sentence_tag[_] == sentence_tag_pred[_] for _ in range(len(sentence))])
        cnt_correct_tag += np.sum([sentence_tag[_] == sentence_tag_pred[_] for _ in range(len(sentence))])
        cnt_total_tag += len(sentence)
        # for j, w in enumerate(sentence):
        #     print('w:{} -> true_tag:{} -> pred_tag:{}'.format(w, sentence_tag[j], sentence_tag_pred[j]))
        # break
    acc = cnt_correct_tag / cnt_total_tag
    # print('cnt_correct_tag, cnt_total_tag:', cnt_correct_tag, cnt_total_tag)
    # print('acc:', acc)
    return acc

def train(Train_Sentences, Train_Sentences_tag):
    '''
    :param Train_Sentences:
    :param Train_Sentences_tag:
    p(Sentences_tag, Sentences) ~  exp(w^T f(Sentences_tag, Sentences)), w是待train的权重，f是特征函数
    :return:
    '''
    N = len(Train_Sentences)
    def get_feature_dict():
        feat_pair2idx = {}
        feat_idx2pair = {}
        feat_num = 0
        for i in range(N):
            sentence = Train_Sentences[i]
            sentence_tag = Train_Sentences_tag[i]
            n = len(sentence)
            pre_tg = 'BOS'
            for i in range(n):
                tg, w = sentence_tag[i], sentence[i]
                if (tg, w) not in feat_pair2idx:
                    feat_pair2idx[(tg, w)] = feat_num
                    feat_idx2pair[feat_num] = (tg, w)
                    feat_num += 1
                if (pre_tg, tg) not in feat_pair2idx:
                    feat_pair2idx[(pre_tg, tg)] = feat_num
                    feat_idx2pair[feat_num] = (pre_tg, tg)
                    feat_num += 1
                pre_tg = tg
            tg = 'EOS'
            if (pre_tg, tg) not in feat_pair2idx:
                feat_pair2idx[(pre_tg, tg)] = feat_num
                feat_idx2pair[feat_num] = (pre_tg, tg)
                feat_num += 1
        return feat_pair2idx, feat_idx2pair, feat_num

    feat_pair2idx, feat_idx2pair, feat_num = get_feature_dict()
    print('{}
{}
{}
'.format(feat_pair2idx, feat_idx2pair, feat_num))

    if reused_W:
        W = np.load(filepath)
    else:
        W = np.random.normal(0, scale=1.0 / np.sqrt(feat_num), size=[feat_num, ])
    # tag, tag2word, tag2tag = {}, {}, {}
    # tag['BOS'] = N
    # tag['EOS'] = N
    Train_Acc, Valid_Acc = [], []
    time_global = time.time()
    for epoch in range(N_epoch):
        time_epoch = time.time()
        s = '###'

        for i in range(N):
            if i % (N // 10)==0:
                s_out = s * (i // (N // 10)) + '{}/{} running this epoch time used: {:.2f}'.format(i, N, time.time() - time_epoch)
                if i // (N // 10) == 10:
                    print(s_out, end="", flush=False)
                else:
                    print(s_out, end="
", flush=True)
            sentence = Train_Sentences[i]
            sentence_tag = Train_Sentences_tag[i]
            n = len(sentence)
            assert len(sentence) == len(sentence_tag)
            assert n > 0
            W_grad = cal_grad_w(W, feat_pair2idx, feat_num, xn=sentence, yn=sentence_tag)
            # W_grad = cal_grad_w_log_version(W, feat_pair2idx, feat_num, xn=sentence, yn=sentence_tag)

            W -= lr * W_grad
        train_acc = evaluate(W, feat_pair2idx, Sentences_=Train_Sentences, Sentences_tag_=Train_Sentences_tag)
        valid_acc = evaluate(W, feat_pair2idx, Sentences_=Valid_Sentences, Sentences_tag_=Valid_Sentences_tag)
        Train_Acc.append(train_acc)
        Valid_Acc.append(valid_acc)
        print('
epoch: {}, epoch time: {}, global time: {}, train acc: {}, valid acc: {}'.format(
            epoch, time.time() - time_epoch, time.time() - time_global, train_acc, valid_acc))

    if show_acc:
        plt.figure()
        plt.title('regulization: {}'.format(regulization))
        plt.plot(Train_Acc, label='Train Acc')
        plt.plot(Valid_Acc, label='Valid Acc')



    return W, feat_pair2idx

REG = [0, 0.1, 0.3, 1, 3, 10, 30]
for reg in REG:
    regulization = reg
    W, feat_pair2idx = train(Train_Sentences, Train_Sentences_tag)
    break
plt.show()
if not os.path.exists(check_path):
    os.makedirs(check_path)
np.save(filepath, W)

# predict_tag(W, feat_pair2idx, sentence=Sentences[0], True_sentence_tag=None)

predict_tag(W, feat_pair2idx, sentence=Sentences[0], True_sentence_tag=Sentences_tag[0])
# predict_tag(sentence=Sentences[0], True_sentence_tag=None)


def fenci_example(W, feat_pair2idx):

    test_sentences = [
        '我是中国人，我爱祖国',
        '独行侠队的球员们承诺每天为达拉斯地区奋战在抗疫一线的工作人员们提供餐食',
        '汤普森太爱打球,不能出场让他很煎熬',
        '这个赛季对克莱来说非常艰难，他太热爱打篮球了，无法上场让他很受打击。',
        '克莱和斯蒂芬会处在极佳的状态，准备好比赛。',
        '勇士已经证明了他们也是一支历史级别的球队，维金斯在稍强于巴恩斯的前提下，仍然算得上是三号位上一位合格的替代者'
    ]

    np.set_printoptions(precision=3, suppress=True)

    for sentence in test_sentences:
        print('-' * 80)
        print('test word : {}'.format(sentence))
        sentence = [w for w in sentence]
        sentence_tag = predict_tag(W, feat_pair2idx, sentence)
        # predict_tag(sentence=Sentences[0], True_sentence_tag=None)
        for i, w in enumerate(sentence):
            print('{} -> {}'.format(w, sentence_tag[i]))

fenci_example(W, feat_pair2idx)

print('end!!!')
'''
test word : 独行侠队的球员们承诺每天为达拉斯地区奋战在抗疫一线的工作人员们提供餐食
独 -> B
行 -> I
侠 -> B
队 -> I
的 -> S
球 -> B
员 -> I
们 -> I
承 -> B
诺 -> I
每 -> B
天 -> I
为 -> B
达 -> I
拉 -> B
斯 -> I
地 -> B
区 -> I
奋 -> B
战 -> I
在 -> S
抗 -> B
疫 -> I
一 -> B
线 -> I
的 -> S
工 -> B
作 -> I
人 -> S
员 -> B
们 -> I
提 -> B
供 -> I
餐 -> B
食 -> I
'''

 　　BILSTM+CRF, PS:为了实现方便，没有加start和end的转移分数权重

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import os
import time
import matplotlib.pyplot as plt
from p03_CRF_layer import CRF_zqx
# from CRF_official import CRF as CRF_zqx

np.random.seed(1)
torch.manual_seed(1)
np.set_printoptions(precision=5, suppress=3)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# device = 'cpu'
print('device:', device)
device = torch.device(device)

small = 50
training = True
show_loss = True
N_epoch = 5
Batch_size = 64
show_epoch = 1
H = 128
em_dim = 100
lr = 1e-2
fold=0
net = 'BiLSTM_CRF'
tag_num = 3

words = {}
max_seq = 0

label2idx = {'B': 0, 'I': 1, 'S': 2}
idx2label = {0: 'B', 1: 'I', 2: 'S'}

Sentences = []
Sentences_label = []
Sentences_origin = []

with open('zhihu.txt', mode='r', encoding='utf8') as f:
    lines = f.readlines()
    print('len(lines):', len(lines))
    for idx, line in enumerate(lines):
        # print(line)
        line = line.split()
        tmp = []
        mmp = []
        for word in line:
            if len(word)==1:
                mmp.append(label2idx['S'])
            else:
                mmp.append(label2idx['B'])
                for _ in range(1, len(word)):
                    mmp.append(label2idx['I'])

            for w in word:
                if w in words:
                    words[w] += 1
                else:
                    words[w] = 1
                tmp.append(w)
        Sentences.append(tmp)
        Sentences_label.append(mmp)
        max_seq = max(max_seq, len(tmp))
        assert len(mmp)==len(tmp)
        # print(tmp)
        if idx > small:
            break


Sentences.sort(key=lambda x: len(x), reverse=True)
Sentences_label.sort(key=lambda x: len(x), reverse=True)
# for sentence in Sentences:
#     print(sentence)
print('len(words):', len(words))
# print(words)
print('max_seq len:', max_seq)

# print(words)
word2index = {word: idx + 1 for idx, word in enumerate(words.keys())}
indx2word = {idx + 1: word for idx, word in enumerate(words.keys())}


voc_size = len(words) + 1
word2index['<pad>'] = 0
indx2word[0] = '<pad>'

Sentences_idx, Sentences_len = [], []
for sentence in Sentences:
    tmp=[]
    for w in sentence:
        tmp.append(word2index[w])
    Sentences_idx.append(torch.LongTensor(tmp))
    Sentences_len.append(len(tmp))
    # print(tmp)
Sentences_idx = torch.nn.utils.rnn.pad_sequence(Sentences_idx,batch_first=True)
# print('-' * 80)
# print(Sentences_idx.size())
# print(Sentences_idx)


# print('-' * 80)
Sentences_label_idx = []
for i, sentences_label in enumerate(Sentences_label):
    tmp = torch.LongTensor(sentences_label)
    Sentences_label_idx.append(tmp)
    # print(Sentences[i])
    # # print(lines[i])
    # print(tmp)
    assert len(tmp) == len(Sentences[i])
Sentences_label_idx = torch.nn.utils.rnn.pad_sequence(Sentences_label_idx,batch_first=True,padding_value=0)
# print('Sentences_label_idx:')
# print(Sentences_label_idx)

# a = torch.tensor(1.0)
# print(a)
# print(a.size())
class MyDataSet(Dataset):
    def __init__(self, data, lens, labels):
        self.data = data
        self.lens = lens
        self.labels = labels
    def __getitem__(self, idx):
        now_data = self.data[idx]
        now_len = self.lens[idx]
        now_mask = []
        now_label = self.labels[idx]
        for i in range(len(now_data)):
            t = 1.0 if i < now_len else 0.0
            now_mask.append(t)
        now_mask = torch.Tensor(now_mask)
        # now_mask = torch.BoolTensor(now_mask)  #用官方CRF的格式要求
        return now_data, now_len, now_mask, now_label
    def __len__(self):
        return len(self.data)

class FenCi_Zqx(nn.Module):
    def __init__(self, voc_size, em_dim, H):
        super(FenCi_Zqx, self).__init__()
        self.emd = nn.Embedding(num_embeddings=voc_size,embedding_dim=em_dim)
        if net == 'LSTM':
            self.rnn = nn.LSTM(input_size=em_dim,hidden_size=H,num_layers=1,batch_first=True)
            self.linear = nn.Linear(in_features=H,out_features=3)
        if 'BiLSTM' in net:
            self.rnn = nn.LSTM(input_size=em_dim, hidden_size=H, num_layers=1, batch_first=True,bidirectional=True)
            self.linear = nn.Linear(in_features=2*H, out_features=3)
        self.loss_fn = CRF_zqx(tag_num=tag_num)
    def forward(self, sentence, sentence_len=None, mask=None):
        emd = self.emd(sentence) #  (batch, seq_len, em_dim)
        all_h, (h, c) = self.rnn(emd) # LSTM: (batch, seq_len, H)   BiLSTM: (batch, seq_len, 2*H)
        # print('emd size:', emd.size())
        # print('all_h.size():', all_h.size())
        # out = all_h.view(-1, all_h.size(2))  # (batch * seq_len, H)
        out = self.linear(all_h).view(emd.size(0), emd.size(1), 3) # (batch, seq_len, 3)
        # print('out size:', out.size())
        return out

Sentences_len = torch.Tensor(Sentences_len)
train_idx = [i for i in range(len(Sentences_len)) if i % 5 == fold]
test_idx = [i for i in range(len(Sentences_len)) if i % 5 != fold]
print(train_idx, '
', test_idx)
Train_Sentences_idx, Train_Sentences_len, Train_Sentences_label_idx = 
    Sentences_idx[train_idx], Sentences_len[train_idx], Sentences_label_idx[train_idx]
Test_Sentences_idx, Test_Sentences_len, Test_Sentences_label_idx = 
    Sentences_idx[test_idx], Sentences_len[test_idx], Sentences_label_idx[test_idx]
Train_data = MyDataSet(Train_Sentences_idx, Train_Sentences_len, Train_Sentences_label_idx)
Train_data_loader = DataLoader(dataset=Train_data, batch_size=Batch_size, shuffle=True)
Test_data = MyDataSet(Test_Sentences_idx, Test_Sentences_len, Test_Sentences_label_idx)
Test_data_loader = DataLoader(dataset=Test_data, batch_size=Batch_size, shuffle=False)

model = FenCi_Zqx(voc_size=voc_size, em_dim=em_dim, H=H)
# loss_fn = nn.CrossEntropyLoss(reduction='none')
# loss_fn = CRF_zqx(tag_num=tag_num)
opt = torch.optim.Adam(model.parameters(), lr=lr)

print('Sentences_idx, Sentences_len, Sentences_label_idx shape')
print(len(Sentences_idx), len(Sentences_len), len(Sentences_label_idx))
print(Sentences_idx.size(), Sentences_len.size(), Sentences_label_idx.size())
print(Sentences_idx.shape, Sentences_len.shape, Sentences_label_idx.shape)
print('#' * 60)
print(model)

def valid(model):
    # model.to(device)
    # model.eval()
    with torch.no_grad():
        avg_loss = 0
        cnt=0
        for batch_step, (now_data, now_len, now_mask, now_label) in enumerate(Test_data_loader):
            cnt += 1
            now_data, now_len, now_mask, now_label = 
                now_data.to(device), now_len.to(device), now_mask.to(device), now_label.to(device)
            out = model(now_data, now_len, now_mask)
            # out = out.view(-1, 3)
            # now_mask = now_mask.view(-1)
            # now_label = now_label.view(-1)

            loss = model.loss_fn(out, now_label, now_mask)
            # print('loss size:', loss.size())
            # print(out.size(), now_label.size(), now_mask.size())
            # loss = torch.mean(loss * now_mask)
            avg_loss += loss.item()
            # print('loss size:', loss.size())
            # print('loss:', loss.item())

        avg_loss /= cnt
        return avg_loss


def train(model):
    print('start training:')
    model.to(device)
    time_st_global = time.time()
    Train_loss,Valid_loss = [], []
    print(model.loss_fn.A)
    # print(model.loss_fn.transitions)
    for epoch in range(N_epoch):
        time_st_epoch = time.time()
        avg_loss = 0
        cnt = 0
        for batch_step, (now_data, now_len, now_mask, now_label) in enumerate(Train_data_loader):
            cnt += 1
            now_data, now_len, now_mask, now_label = 
                now_data.to(device), now_len.to(device), now_mask.to(device), now_label.to(device)
            out = model(now_data, now_len, now_mask)
            # out = out.view(-1, 3)
            # now_mask = now_mask.view(-1)
            # now_label = now_label.view(-1)
            loss = model.loss_fn(out, now_label, now_mask)
            # print('loss size:', loss.size())
            # print(out.size(), now_label.size(), now_mask.size())
            # loss = torch.mean(loss * now_mask)
            avg_loss += loss.item()
            # print('loss size:', loss.size())
            # print('loss:', loss.item())

            opt.zero_grad()
            loss.backward()
            opt.step()
        avg_loss /= cnt
        valid_avg_loss = valid(model)
        print('#' * 80)
        print('epoch:{}, steps: {}, train avg loss: {} -- valid avg loss : {} '.format(epoch, cnt, avg_loss, valid_avg_loss))
        print('epoch time {:.2f} s, total time {:.2f} s'.format(time.time() - time_st_epoch,
                                                                time.time() - time_st_global))
        if len(Valid_loss)==0 or valid_avg_loss < min(Valid_loss):
            if not os.path.exists(check_path):
                os.makedirs(check_path)
            torch.save(model.state_dict(), filepath)

        Train_loss.append(avg_loss)
        Valid_loss.append(valid_avg_loss)

    print(model.loss_fn.A)
    # print(model.loss_fn.transitions)
    if show_loss:
        plt.figure()
        plt.plot(Train_loss,label='Train loss')
        plt.plot(Valid_loss, label='Valid loss')
        plt.legend()
        plt.savefig('Train_Valid_loss' + net + '.png')
        # plt.show()

    return model

        # break

check_path = './Checkpoints/'
filepath = check_path + 'p03_Fenci_state_dict_' + net + ' .pkl'

if training:
    model = train(model)


# 模型恢复
model.load_state_dict(torch.load(filepath))


test_words = [
    '我是中国人，我爱祖国',
    '独行侠队的球员们承诺每天为达拉斯地区奋战在抗疫一线的工作人员们提供餐食',
    '汤普森太爱打球,不能出场让他很煎熬',
    '这个赛季对克莱来说非常艰难，他太热爱打篮球了，无法上场让他很受打击。',
    '克莱和斯蒂芬会处在极佳的状态，准备好比赛。',
    '勇士已经证明了他们也是一支历史级别的球队，维金斯在稍强于巴恩斯的前提下，仍然算得上是三号位上一位合格的替代者'
]

np.set_printoptions(precision=3, suppress=True)
model.cpu()
model.eval()
for word in test_words:
    print('-' * 80)
    print('test word : {}'.format(word))
    word_idx = [word2index[w] for w in word]
    word_idx = torch.LongTensor([word_idx])
    # print('word_idx.size():', word_idx.size())
    # word_idx.to(device)
    out = model(word_idx)

    # out = model.loss_fn.decode(emissions=out, mask=None)
    out = model.loss_fn.decode(y_pred=out, mask=None)
    out_label = out[0]


    # # print('out.size():', out.size())
    # out = out.squeeze(0).cpu().detach().numpy()
    # # print('out.shape():', out.shape)
    #
    # out_label = np.argmax(out, axis=1)
    # # print(out_label)

    for i, w in enumerate(word):
        print('{} -> {} -> {}'.format(w, idx2label[out_label[i]], out_label[i]))

print('end!!!')
'''
test word : 勇士已经证明了他们也是一支历史级别的球队，维金斯在稍强于巴恩斯的前提下，仍然算得上是三号位上一位合格的替代者
勇 -> B -> 0
士 -> I -> 1
已 -> B -> 0
经 -> I -> 1
证 -> B -> 0
明 -> I -> 1
了 -> S -> 2
他 -> B -> 0
们 -> I -> 1
也 -> S -> 2
是 -> S -> 2
一 -> S -> 2
支 -> S -> 2
历 -> B -> 0
史 -> I -> 1
级 -> B -> 0
别 -> I -> 1
的 -> S -> 2
球 -> B -> 0
队 -> I -> 1
， -> S -> 2
维 -> B -> 0
金 -> I -> 1
斯 -> I -> 1
在 -> S -> 2
稍 -> B -> 0
强 -> I -> 1
于 -> S -> 2
巴 -> B -> 0
恩 -> I -> 1
斯 -> I -> 1
的 -> S -> 2
前 -> B -> 0
提 -> I -> 1
下 -> S -> 2
， -> S -> 2
仍 -> B -> 0
然 -> I -> 1
算 -> S -> 2
得 -> S -> 2
上 -> S -> 2
是 -> S -> 2
三 -> S -> 2
号 -> S -> 2
位 -> S -> 2
上 -> S -> 2
一 -> S -> 2
位 -> S -> 2
合 -> B -> 0
格 -> I -> 1
的 -> S -> 2
替 -> B -> 0
代 -> I -> 1
者 -> I -> 1


'''

　　附上自己写的CRF模块以及公式注解

import torch
import torch.nn as nn
import numpy as np
np.random.seed(1)
torch.manual_seed(1)

class CRF_zqx(nn.Module):
    def __init__(self, tag_num):
        super(CRF_zqx, self).__init__()
        # A为转移矩阵, A_ij, 表示tag i 到 tag j 的得分
        # self.A = torch.rand(size=(tag_num, tag_num), requires_grad=True)
        # self.A = nn.Parameter(torch.rand(size=(tag_num, tag_num)))
        self.A = nn.Parameter(torch.empty(tag_num, tag_num))
        self.tag_num = tag_num
        self.reset_parameters()
    def reset_parameters(self) -> None:
        """Initialize the transition parameters.

        The parameters will be initialized randomly from a uniform distribution
        between -0.1 and 0.1.
        """
        nn.init.uniform_(self.A, -0.1, 0.1)

    def forward(self, y_pred, y_true, mask):
        if len(y_true.size()) < 3:
            # print(y_true.dtype)
            y_true = torch.nn.functional.one_hot(y_true, num_classes=self.tag_num)
            y_true = y_true.type(torch.float32)
        # y_pred, y_true: [batch_size, seq_len, tag_num],   ps:y_true是one-hot向量
        # log p(y_true | x_true) = log {exp(score(y_true, x_true) / Sigma_y exp(score(y, x_true))}
        #                        = score(y_true, x_true) - log  sum_y exp(score(y, x_true))
        # print('forward:
')
        # print('y_pred:{}
y_true:{}
mask:{}
'.format(y_pred, y_true, mask))
        # print('y_pred:{}
y_true:{}
mask:{}
'.format(y_pred.size(), y_true.size(), mask.size()))
        # print('A:', self.A)
        loss = self.score(y_pred, y_true, mask) - self.log_sum_exp(y_pred, mask)
        return torch.mean(-loss)

    def score(self, y_pred, y_true, mask):
        # y_pred, y_true: [batch_size, seq_len, tag_num]  mask: [batch_size, seq_len]
        mask = torch.unsqueeze(mask, dim=2)  #  mask: [batch_size, seq_len, 1]
        # print('y_pred, y_true, mask size:', y_pred.size(), y_true.size(), mask.size())
        score_word2tag = torch.sum(y_pred * y_true * mask, dim=[1, 2])  #  计算word2tag的分数,得到[batch_size, ]向量
        #            [batch_size, seq_len-1, tag_num, 1]  *  [batch_size, seq_len-1, 1, tag_num]
        #    从而获得[batch_size, seq_len-1, tag_num, tag_num], 后两个维度都是one-hot向量，分别表示tag2tag的转移矩阵A的index
        score_tag2tag = torch.unsqueeze(y_true[:, :-1, :] * mask[:, :-1, :], dim=3) 
                        * torch.unsqueeze(y_true[:, 1:, :] * mask[:, 1:, :], dim=2)

        #               [batch_size, seq_len-1, tag_num, tag_num]  *  [1, 1, tag_num, tag_num]
        A = torch.unsqueeze(torch.unsqueeze(self.A, 0), 0)
        score_tag2tag = score_tag2tag * A
        score_tag2tag = torch.sum(score_tag2tag, dim=[1, 2, 3])  # [batch_size,]
        score_ = score_word2tag + score_tag2tag
        # print('score_ size:', score_.size())
        # print('score:', score_)
        return score_

    def log_sum_exp(self, y_pred, mask):
        # mask: [batch_size, seq_len]
        seq_len = y_pred.size(1)
        pre_log_Z = y_pred[:, 0, :] # [batch_size, tag_num], initial: log Z = log exp(y_pred[time_step=0]) = y_pred[:, 0 , :]

        # print('pre_log_Z:{}, with size:{}'.format(pre_log_Z, pre_log_Z.size()))
        for i in range(1, seq_len):
            # print('i:', i)
            #                    [1, tag_num, tag_num]   +  [batch_size, tag_num, 1] = [batch_size, tag_num, tag_num]
            # 然后对列(dim=1)求logsumexp,  得到[batch_size, tag_num]
            tmp = pre_log_Z.unsqueeze(2)
            # log_Z = torch.logsumexp(tmp + self.A + y_pred[:, i:i+1, :], dim=1)
            log_Z = torch.logsumexp(torch.unsqueeze(self.A, 0) + torch.unsqueeze(pre_log_Z, 2), dim=1) + y_pred[:, i, :]
            log_Z = mask[:, i:i+1] * log_Z + (1 - mask[:, i:i+1]) * pre_log_Z  # 现在mask位置上是1，则更新， 如果是0，则取用pre_log_Z的值
            pre_log_Z = log_Z.clone()
        # print('log_Z size:', pre_log_Z.size())

        # print('res:', pre_log_Z)
        res = torch.logsumexp(pre_log_Z,dim=1)  # 是logsumexp  不是 sum,  debug了大半天！！！！
        # print('logsumexp:', res)
        return res

    def decode(self,y_pred, mask=None):
        batch, seq_len = y_pred.size(0), y_pred.size(1)
        if mask is None:
            mask = torch.ones(size=[batch, seq_len])

        pre_dp = y_pred[:, 0, :]  #[batch, tag_num]
        dp_best_idx = torch.LongTensor(torch.zeros(size=[batch, seq_len + 1, self.tag_num], dtype=torch.long) - 1)
        for i in range(1, seq_len):                      # from     to
            now_pred = y_pred[:, i:i+1, :]       # [batch, 1,       tag_num]
            pre_dp = torch.unsqueeze(pre_dp, 2)  # [batch, tag_num, 1      ]
            A = torch.unsqueeze(self.A, 0)       # [1,     tag_num, tag_num]
            dp, idx = torch.max(pre_dp + A + now_pred, dim=1) #  dp: [batch, tag_num]
            # print('dp:{}, idx:{}'.format(dp.size(), idx.size()))
            dp_best_idx[:, i, :] = idx
            pre_dp = dp.clone()

        best_value, last_tag = torch.max(pre_dp, dim=1)
        print('pre_dp:{}, pre_dp size:{}
pointer:{}, last_tag size:{}'.format(pre_dp, pre_dp.size(), last_tag, last_tag.size()))
        last_tag = list(last_tag.cpu().detach().numpy())
        dp_best_idx = dp_best_idx.cpu().detach().numpy()
        print('last tag:', last_tag)
        ans = [last_tag] # [batch]
        i = seq_len - 1
        while i:
            tmp = dp_best_idx[:, i, :]
            pre_tag = []
            for j in range(batch):
                pre_tag.append(tmp[j, last_tag[j]])
            last_tag = pre_tag.copy()
            ans = [pre_tag] + ans
            i -= 1
        ans = np.array(ans) #[seq_len, batch]
        ans = ans.transpose()
        print('ans:', ans)
        # while i:
        #     print('dp_best_idx[:, i, :] size:{}, pointer.unsqueeze(1) size:{}'.format(
        #         dp_best_idx[:, i, :].size(), pointer.unsqueeze(1).size()))
        #     print('dp_best_idx[:, i, :]:{}, pointer.unsqueeze(1):{}'.format(
        #         dp_best_idx[:, i, :], pointer.unsqueeze(1)))
        #     pointer = dp_best_idx[:, i, :][pointer.unsqueeze(1)]  # pointer.unsqueeze(1): [batch, 1]
        #     ans = [list(pointer)] + ans
        #     i = i - 1

        return ans

if __name__=='__main__':
    batch = 1
    seq_len = 3
    tag_num = 2
    y_pred = torch.rand(size=[batch, seq_len, tag_num])
    y_true = torch.randint(0, tag_num, size=[batch, seq_len])
    # print(y_true)
    y_true = torch.nn.functional.one_hot(y_true, num_classes=tag_num)
    y_true = y_true.type(torch.float32)
    # print(y_true)
    # print(y_true.size())
    mask = []
    for _ in range(batch):
        tmp = np.random.randint(2, seq_len)
        mask.append([1] * tmp + [0] * (seq_len - tmp))
    mask = torch.Tensor(mask)
    # print(mask)
    model =CRF_zqx(tag_num=tag_num)



    # print('y_pred:{}
y_true:{}
mask:{}
'.format(y_pred, y_true, mask))
    # print(type(y_pred))
    # print(type(y_true))
    # print(type(mask))
    # print(y_pred.dtype)
    # print(y_true.dtype)
    # print(mask.dtype)

    print('y_pred=y_pred, y_true=y_true, mask=mask:', y_pred.size(), y_true.size(), mask.size())
    loss = model(y_pred=y_pred, y_true=y_true, mask=mask)
    print('loss: {}'.format(loss))


'''
y_pred:tensor([[[0.7576, 0.2793],
         [0.4031, 0.7347]]])
y_true:tensor([[[0., 1.],
         [0., 1.]]])
mask:tensor([[1., 0.]])

'''

LDA 模型（具体请看LDA数学八卦）