def train(epoch):
print('\nEpoch: %d' % epoch)
global Train_acc
net.train()
train_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(trainloader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
optimizer.zero_grad()
#inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
utils.clip_gradient(
optimizer, 0.1) #clip_gradient能有效了控制梯度爆炸的影响，使得最终的loss能下降到满意的结果
optimizer.step()
train_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
utils.progress_bar(
batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
Train_acc = 100. * correct / total
def test(epoch):
global Test_acc
global best_Test_acc
global best_Test_acc_epoch
net.eval()
PrivateTest_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(testloader):
bs, ncrops, c, h, w = np.shape(inputs)
inputs = inputs.view(-1, c, h, w)
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
outputs_avg = outputs.view(bs, ncrops, -1).mean(1) # avg over crops
loss = criterion(outputs_avg, targets)
PrivateTest_loss += loss.item()
_, predicted = torch.max(outputs_avg.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
utils.progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (PrivateTest_loss / (batch_idx + 1), \
100. * correct / total, correct, total))
# Save checkpoint.
Test_acc = 100. * correct / total
if Test_acc > best_Test_acc:
print('Saving..')
print("best_Test_acc: %0.3f" % Test_acc)
state = {
'net': net.state_dict() if use_cuda else net,
'best_Test_acc': Test_acc,
'best_Test_acc_epoch': epoch,
}
if not os.path.isdir(opt.dataset + '_' + 'ShuffleNetV2'):
os.mkdir(opt.dataset + '_' + 'ShuffleNetV2')
if not os.path.isdir(path):
os.mkdir(path)
#torch.save(state, os.path.join(path, str(best_Test_acc) + '_ShuffleNetV2.pth'))
best_Test_acc = Test_acc
best_Test_acc_epoch = epoch
torch.save(
state,
os.path.join(
path,
str(best_Test_acc) + '_' + str(input_size) +
'_ShuffleNetV2.pth'))
def train(epoch):
print('\nEpoch: %d' % epoch)
global Train_acc
net.train()
train_loss = 0
correct = 0
total = 0
'''
if epoch > learning_rate_decay_start and learning_rate_decay_start >= 0:
frac = (epoch - learning_rate_decay_start) // learning_rate_decay_every
decay_factor = learning_rate_decay_rate ** frac
current_lr = opt.lr * decay_factor
utils.set_lr(optimizer, current_lr) # set the decayed rate
else:
current_lr = opt.lr
print('learning_rate: %s' % str(current_lr))
'''
for batch_idx, (inputs, targets) in enumerate(trainloader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
optimizer.zero_grad()
inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
utils.clip_gradient(optimizer, 0.1)
optimizer.step()
train_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
utils.progress_bar(
batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
Train_acc = 100. * correct / total
def train(self, nb_epochs, data_loader):
"""Train model on data"""
# Initialize tracked quantities and prepare everything
G_all_losses, D_all_losses, times = [], [], utils.AvgMeter()
utils.format_hdr(self.gan, self.root_dir, self.train_len)
start = datetime.datetime.now()
g_iter, d_iter = 0, 0
# Train
for epoch in range(nb_epochs):
print('EPOCH {:d} / {:d}'.format(epoch + 1, nb_epochs))
G_losses, D_losses = utils.AvgMeter(), utils.AvgMeter()
start_epoch = datetime.datetime.now()
avg_time_per_batch = utils.AvgMeter()
# Mini-batch SGD
for batch_idx, (x, _) in enumerate(data_loader):
# Critic update ratio
if self.gan_type == 'wgan':
n_critic = 20 if g_iter < 50 or (
g_iter + 1) % 500 == 0 else self.n_critic
else:
n_critic = self.n_critic
# Training mode
self.gan.G.train()
# Discard last examples to simplify code
if x.size(0) != self.batch_size:
break
batch_start = datetime.datetime.now()
# Print progress bar
utils.progress_bar(batch_idx, self.batch_report_interval,
G_losses.avg, D_losses.avg)
x = Variable(x)
if torch.cuda.is_available() and self.use_cuda:
x = x.cuda()
# Update discriminator
D_loss, fake_imgs = self.gan.train_D(x, self.D_optimizer,
self.batch_size)
D_losses.update(D_loss, self.batch_size)
d_iter += 1
# Update generator
if batch_idx % n_critic == 0:
G_loss = self.gan.train_G(self.G_optimizer,
self.batch_size)
G_losses.update(G_loss, self.batch_size)
g_iter += 1
batch_end = datetime.datetime.now()
batch_time = int(
(batch_end - batch_start).total_seconds() * 1000)
avg_time_per_batch.update(batch_time)
# Report model statistics
if (batch_idx % self.batch_report_interval == 0 and batch_idx) or \
self.batch_report_interval == self.num_batches:
G_all_losses.append(G_losses.avg)
D_all_losses.append(D_losses.avg)
utils.show_learning_stats(batch_idx, self.num_batches,
G_losses.avg, D_losses.avg,
avg_time_per_batch.avg)
[
k.reset()
for k in [G_losses, D_losses, avg_time_per_batch]
]
self.eval(100, epoch=epoch, while_training=True)
# print('Critic iter: {}'.format(g_iter))
# Save stats
if batch_idx % self.save_stats_interval == 0 and batch_idx:
stats = dict(G_loss=G_all_losses, D_loss=D_all_losses)
self.save_stats(stats)
# Save model
utils.clear_line()
print('Elapsed time for epoch: {}'.format(
utils.time_elapsed_since(start_epoch)))
self.gan.save_model(self.ckpt_path, epoch)
self.eval(100, epoch=epoch, while_training=True)
# Print elapsed time
elapsed = utils.time_elapsed_since(start)
print('Training done! Total elapsed time: {}\n'.format(elapsed))
return G_loss, D_loss