def start_train():
'''
训练
'''
use_amp = True
# 前向反传N次，再更新参数 目的：增大batch（理论batch= batch_size * N）
iter_size = 8
myNet = MyNet(use_amp).to("cuda:0")
myNet = torch.nn.DataParallel(myNet, device_ids=[0, 1]) # 数据并行
myNet.train()
# 训练开始前初始化 梯度缩放器
scaler = GradScaler() if use_amp else None
# 加载预训练权重
if resume_train:
scaler.load_state_dict(checkpoint['scaler']) # amp自动混合精度用到
optimizer.load_state_dict(checkpoint['optimizer'])
myNet.load_state_dict(checkpoint["model"])
for epoch in range(1, 100):
for batch_idx, (input, target) in enumerate(dataloader_train):
# 数据 转到每个并行模型的主卡上
input = input.to("cuda:0")
target = target.to("cuda:0")
# 自动混合精度训练
if use_amp:
# 自动广播 将支持半精度操作自动转为FP16
with autocast():
# 提取特征
feature = myNet(input)
losses = loss_function(target, feature)
loss = losses / iter_size
scaler.scale(loss).backward()
else:
feature = myNet(input, target)
losses = loss_function(target, feature)
loss = losses / iter_size
loss.backward()
# 梯度累积,再更新参数
if (batch_idx + 1) % iter_size == 0:
# 梯度更新
if use_amp:
scaler.step(optimizer)
scaler.update()
else:
optimizer.step()
# 梯度清零
optimizer.zero_grad()
# scaler 具有状态。恢复训练时需要加载
state = {
'net': myNet.state_dict(),
'optimizer': optimizer.state_dict(),
'scaler': scaler.state_dict()
}
torch.save(state, "filename.pth")
input_signature = input_signature.to(device)
model = model.to(device)
pruning_runner = get_pruning_runner(model, input_signature, 'iterative')
model = pruning_runner.prune(removal_ratio=args.sparsity, mode='sparse')
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, args.epochs)
best_acc1 = 0
for epoch in range(args.epochs):
train(train_loader, model, criterion, optimizer, epoch)
lr_scheduler.step()
acc1, acc5 = evaluate(val_loader, model, criterion)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if is_best:
if hasattr(model, 'module'):
torch.save(model.state_dict(), model_path)
else:
torch.save(model.state_dict(), model_path)
download = True
train_loader = get_dataloader_ddp(args.data_dir,
batch_size,
num_workers=args.num_workers,
shuffle=True,
train=True,
download=download)
val_loader = get_dataloader_ddp(args.data_dir,
batch_size,
num_workers=args.num_workers,
shuffle=False,
train=False,
download=download)
criterion = torch.nn.CrossEntropyLoss().cuda()
if not os.path.exists(args.pretrained):
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, args.epochs)
best_acc1 = 0
for epoch in range(args.epochs):
train(train_loader, model, criterion, optimizer, epoch)
lr_scheduler.step()
acc1, acc5 = evaluate(val_loader, model, criterion)
if acc1 > best_acc1:
best_acc1 = acc1
torch.save(model.state_dict(), args.pretrained)
def do_train(data_path,
model_name='mymodel',
use_gpu=False,
epoch_num=5,
batch_size=100,
learning_rate=0.01):
place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
with fluid.dygraph.guard(place):
model = MyNet()
model.train()
train_loader = load_data(data_path, mode='train')
optimizer = fluid.optimizer.SGDOptimizer(
learning_rate=learning_rate, parameter_list=model.parameters())
iter = 0
for epoch_id in range(epoch_num):
for batch_id, data in enumerate(train_loader()):
#准备数据，格式需要转换成符合框架要求的
image_data, label_data = data
# 将数据转为飞桨动态图格式
image = fluid.dygraph.to_variable(image_data)
label = fluid.dygraph.to_variable(label_data)
# #前向计算的过程
# predict = model(image)
#前向计算的过程，同时拿到模型输出值和分类准确率
predict, avg_acc = model(image, label)
#计算损失，取一个批次样本损失的平均值
# loss = fluid.layers.square_error_cost(predict, label)
loss = fluid.layers.cross_entropy(predict, label)
avg_loss = fluid.layers.mean(loss)
#每训练了1000批次的数据，打印下当前Loss的情况
if batch_id != 0 and batch_id % 100 == 0:
print(
"epoch: {}, batch: {}, loss is: {}, acc is: {}".format(
epoch_id, batch_id, avg_loss.numpy(),
avg_acc.numpy()))
log_writer.add_scalar(tag='acc',
step=iter,
value=avg_acc.numpy())
log_writer.add_scalar(tag='loss',
step=iter,
value=avg_loss.numpy())
iter = iter + 100
#后向传播，更新参数的过程
avg_loss.backward()
optimizer.minimize(avg_loss)
model.clear_gradients()
fluid.save_dygraph(
model.state_dict(),
os.path.join(CHECKPOINT_PATH,
f'{model_name}_epoch_{epoch_id}'))
fluid.save_dygraph(
optimizer.state_dict(),
os.path.join(CHECKPOINT_PATH,
f'{model_name}_epoch_{epoch_id}'))
# 保存模型
fluid.save_dygraph(model.state_dict(),
os.path.join(MODEL_PATH, model_name))
mode='slim',
index=args.subnet_index)
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, args.epochs)
best_acc1 = 0
for epoch in range(args.epochs):
train(train_loader, model, criterion, optimizer, epoch)
lr_scheduler.step()
acc1, acc5 = evaluate(val_loader, model, criterion)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if is_best:
if hasattr(model, 'module'):
torch.save(model.state_dict(),
os.path.join(args.save_dir, 'mynet_slim.pth'))
else:
torch.save(model.state_dict(),
os.path.join(args.save_dir, 'mynet_slim.pth'))
# NOTE: Check out the torch.optim library for other optimization algorithms that
# could be used instead.
optimizer = optim.SGD(net.parameters(), lr=learning_rate)
# Begin the training loop.
for epoch in range(num_epoch):
# Helps keep track of where you are in training. Don't like this method?
# Check out the tqdm module, which will print a loading bar for for loops.
print("Epoch:", epoch + 1)
for data in trainloader:
# Get the inputs with labels
inputs, labels = data
# Zero the parameter gradients
optimizer.zero_grad()
# Feedforward
#inputs = inputs.view(-1, 784)
outputs = net(inputs)
# Backpropogation
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# Save the trained neural network.
torch.save(net.state_dict(), "./" + net_name + ".pth")
print("Finished training! Saved model parameters in file " + net_name + ".pth")
class Trainer:
def __init__(self, net_path, board_size=15, n=5):
# 游戏棋盘大小
self.board_size = board_size
# 连五子胜利
self.n = n
# 环境实例化
self.env = Game(board_size, n)
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.number_playout = args.n_playout
# 记忆库大小
self.buffer_size = args.buffer_size
self.buffer = deque(maxlen=self.buffer_size)
self.batch_size = args.batch_size
# 自我对局1次后进行训练
self.n_games = args.n_games
# 自我对局后进行5次训练
self.epochs = args.epochs
# 打印保存模型间隔
self.check_freq = args.check_freq
# 总共游戏次数
self.game_num = args.game_num
self.net_path = net_path
self.net = MyNet().to(self.device)
self.MSELoss = nn.MSELoss()
self.optimizer = torch.optim.Adam(self.net.parameters(),
weight_decay=1e-4)
# 实例化蒙特卡洛玩家，参数：游戏策略，探索常数，模拟次数，是否自我对弈（测试时为False）
self.mcts_player = Player(policy=self.policy,
number_playout=self.number_playout,
is_self_play=True)
self.writer = SummaryWriter()
if os.path.exists(net_path):
self.net.load_state_dict(torch.load(net_path))
else:
self.net.apply(self.weight_init)
def weight_init(self, net):
if isinstance(net, nn.Linear) or isinstance(net, nn.Conv2d):
nn.init.normal_(net.weight, mean=0., std=0.1)
nn.init.constant_(net.bias, 0.)
def train(self):
for i in range(self.game_num):
# 环境先自我对弈获得棋局状态，动作概率以及玩家可以赢的概率值
for _ in range(self.n_games):
winner, data = self.env.self_play(self.mcts_player, temp=1.0)
# 打印每局对局信息
print(self.env, "\n", "------------------xx--------")
# 将获得的数据多样化存入样本池
self.extend_sample(data)
# 取样训练
batch = random.sample(self.buffer,
min(len(self.buffer), self.batch_size))
# 解包
state_batch, mcts_probs_batch, winner_value_batch = zip(*batch)
loss = 0.
for _ in range(self.epochs):
# 数据处理
state_batch = torch.tensor(state_batch).to(self.device)
mcts_probs_batch = torch.tensor(mcts_probs_batch).to(
self.device)
winner_value_batch = torch.tensor(winner_value_batch).to(
self.device)
# 通过神经网络输出动作概率，价值用于训练
log_act_probs, value = self.net(state_batch)
# 计算损失
# 价值损失：输出价值与该状态所在对局最终胜负的值（-1/0/1）（均方差）
# 策略损失：蒙特卡洛树模拟的概率值与神经网络模拟的概率值的相似度 (-log(pi) * p)(交叉熵)
value_loss = self.MSELoss(value,
winner_value_batch.view_as(value))
policy_loss = -torch.mean(
torch.sum(mcts_probs_batch * log_act_probs, dim=-1))
loss = value_loss + policy_loss
# 反向传播
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
print(f"epoch:{i},loss:{loss}")
self.writer.add_scalar("loss", loss, i)
self.net.add_histogram(self.writer, i)
if (i + 1) % self.check_freq == 0:
torch.save(self.net.state_dict(), self.net_path)
# 多样化数据样本
def extend_sample(self, data):
extend_data = []
for state, mcts_prob, winner_value in data:
extend_data.append((state, mcts_prob, winner_value))
# 分别旋转 90度/180度/270度，增加数据多样性
for i in range(1, 4):
# 同时旋转棋盘状态和概率值
state_ = np.rot90(state, i, (1, 2))
mcts_prob_ = np.rot90(
mcts_prob.reshape(self.env.height, self.env.width), i)
extend_data.append(
(state_, mcts_prob_.flatten(), winner_value))
# 翻转棋盘,将矩阵中的每一位玩家的状态进行翻转
state_ = np.array([np.fliplr(s) for s in state_])
mcts_prob_ = np.fliplr(mcts_prob_)
extend_data.append(
(state_, mcts_prob_.flatten(), winner_value))
# 将样本存入样本池
self.buffer.extend(extend_data)
# 用于player调用神经网络获得动作概率，当前局面价值
def policy(self, env):
# 获取可用动作 15*15=225
action_avail = env.action_avail
# 获得当前状态
state = torch.from_numpy(env.get_state).unsqueeze(0).to(self.device)
# 放入神经网络得到预测的log动作概率以及当前状态的胜率
log_action_probs, value = self.net(state)
# 把 log 动作概率转换为动作概率并过滤不可用动作
act_probs = torch.exp(
log_action_probs).detach().cpu().numpy().flatten()
act_probs = zip(action_avail, act_probs[action_avail])
value = value.item()
# 返回动作概率，当前局面价值
return act_probs, value