-
Notifications
You must be signed in to change notification settings - Fork 750
/
训练_B.py
138 lines (107 loc) · 5.05 KB
/
训练_B.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import torch
import torchvision
from PIL import Image
import numpy as np
import time
import json
from config import GPT2Config, TransformerConfig
from Batch import create_masks
from ModelA import get_model
import torch.nn.functional as F
from 取训练数据 import *
from 杂项 import *
import os
import random
训练数据保存目录='../训练数据样本'
if not os.path.exists(训练数据保存目录):
os.makedirs(训练数据保存目录)
for root, dirs, files in os.walk('../训练数据样本'):
if len(dirs)>0:
break
词数词典路径="./json/词_数表.json"
数_词表路径="./json/数_词表.json"
if os.path.isfile(词数词典路径) and os.path.isfile(数_词表路径):
词_数表, 数_词表 = 读出引索(词数词典路径, 数_词表路径)
with open(词数词典路径, encoding='utf8') as f:
词数词典=json.load(f)
device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
#
#
config = TransformerConfig()
model = get_model(config, 130)
模型路径 = 'weights/model_weights'
model = model.cuda(device)
optimizer = torch.optim.Adam(model.parameters(), lr=6.25e-5, betas=(0.9, 0.98), eps=1e-9)
分块大小=25
游标大小=23
树枝=10
计数=0
time_start=time.time()
for j in range(100):
random.shuffle(dirs)
for 号 in dirs:
预处理数据 = '../训练数据样本/'+号+'/图片_操作预处理数据2.npz'
if os.path.isfile(预处理数据):
npz文件 = np.load(预处理数据, allow_pickle=True)
图片张量np, 操作序列 = npz文件["图片张量np"], npz文件["操作序列"]
循环=True
游标=0
操作序列=np.insert(操作序列,0,128)
操作_分_表 = []
目标输出_分_表 = []
图片_分_表 = []
while 循环:
if 游标 + 分块大小 < 操作序列.shape[0]:
操作_分 = 操作序列[游标:游标 + 分块大小]
目标输出_分 = 操作序列[游标 + 1:游标 + 1 + 分块大小]
图片_分 = 图片张量np[游标:游标 + 分块大小, :]
操作_分_表.append(操作_分)
目标输出_分_表.append(目标输出_分)
图片_分_表.append(图片_分)
游标 = 游标 + 游标大小
else:
操作_分 = 操作序列[-分块大小 - 1:-1]
目标输出_分 = 操作序列[-分块大小:]
图片_分 = 图片张量np[-分块大小:, :]
操作_分_表.append(操作_分)
目标输出_分_表.append(目标输出_分)
图片_分_表.append(图片_分)
循环 = False
循环=True
i=0
while 循环:
if (i+1)*树枝<len(操作_分_表):
操作_分_枝=np.array(操作_分_表[i*树枝:(i+1)*树枝])
图片_分_枝 = np.array(图片_分_表[i * 树枝:(i + 1) * 树枝])
目标输出_分_枝 = np.array(目标输出_分_表[i * 树枝:(i + 1) * 树枝])
else:
操作_分_枝 = np.array(操作_分_表[i * 树枝:len(操作_分_表)])
图片_分_枝 = np.array(图片_分_表[i * 树枝:len(图片_分_表)],dtype=np.float32)
目标输出_分_枝 = np.array(目标输出_分_表[i * 树枝:len(目标输出_分_表)])
循环 = False
操作_分_torch=torch.from_numpy(操作_分_枝).cuda(device)
图片_分_torch = torch.from_numpy(图片_分_枝).cuda(device)
目标输出_分_torch = torch.from_numpy(目标输出_分_枝).cuda(device)
src_mask, trg_mask = create_masks(操作_分_torch, 操作_分_torch, device)
if 图片_分_torch.shape[0]!=操作_分_torch.shape[0]:
continue
输出_实际_A = model(图片_分_torch,操作_分_torch ,trg_mask)
lin = 输出_实际_A.view(-1, 输出_实际_A.size(-1))
optimizer.zero_grad()
loss = F.cross_entropy(lin, 目标输出_分_torch.contiguous().view(-1), ignore_index=-1)
if 计数 % 1 == 0:
print(loss)
time_end = time.time()
用时 = time_end - time_start
_, 抽样 = torch.topk(输出_实际_A, k=1, dim=-1)
抽样np = 抽样.cpu().numpy()
打印抽样数据(数_词表, 抽样np[0:1,:,:], 目标输出_分_torch[0,:])
print("用时{} 第{}轮 第{}张 号{}".format(用时, j, 计数, 号))
if 计数 % 45060 == 0:
print('888')
loss.backward()
optimizer.step()
计数=计数+1
i=i+1
torch.save(model.state_dict(), 'weights/model_weights')
torch.save(model.state_dict(), 'weights/model_weights_P{}'.format(str(j)))