def train(epoch):
net.train()
train_loss = 0
correct = 0
total = 0
p = None
for batch_idx, (inputs, targets) in enumerate(trainloader):
x_stack = None
imgs = inputs.numpy().astype('float32')
# block scrambling
x_stack = blockwise_scramble(imgs)
imgs = np.transpose(x_stack, (0, 3, 1, 2))
inputs = torch.from_numpy(imgs)
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = net(inputs)
true_loss = criterion(outputs, targets)
loss = true_loss
loss.backward()
optimizer.step()
train_loss += true_loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
return train_loss, 100. * correct / total
def train(epoch):
net.train()
train_loss = 0
correct = 0
total = 0
p = None
param12 = 0.001
param3 = 0.001
param4 = 1e-1
for batch_idx, (inputs, targets) in enumerate(trainloader):
x_stack = None
imgs = inputs.numpy().astype('float32')
# block scrambling
x_stack = blockwise_scramble(imgs)
imgs = np.transpose(x_stack, (0, 3, 1, 2))
inputs = torch.from_numpy(imgs)
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs, mat, feature = net(inputs)
true_loss = criterion(outputs, targets)
# doubly stochastic constraint
dsc = 0
for i in range(64):
dsc += torch.abs(mat[i, :]).sum() - torch.sqrt(
(mat[i, :] * mat[i, :]).sum())
dsc += torch.abs(mat[:, i]).sum() - torch.sqrt(
(mat[:, i] * mat[:, i]).sum())
dsc = param3 * dsc / (64 * 64)
natural_image_prior = param4 * total_variation_norm(
feature) / inputs.size()[0]
loss = true_loss + dsc + natural_image_prior
loss.backward()
optimizer.step()
train_loss += true_loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
return train_loss, 100. * correct / total
def test(epoch):
global best_acc
net.eval()
test_loss = 0
correct = 0
total = 0
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(testloader):
x_stack = None
imgs = inputs.numpy().astype('float32')
# block scrambling
x_stack = blockwise_scramble(imgs)
imgs = np.transpose(x_stack, (0, 3, 1, 2))
inputs = torch.from_numpy(imgs)
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = net(inputs)
true_loss = criterion(outputs, targets)
test_loss += true_loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
# Save checkpoint.
acc = 100. * correct / total
if best_acc < acc:
# print('Saving..')
state = {
'net': net.state_dict(),
'acc': acc,
'epoch': epoch,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, './' + args.training_model_name)
best_acc = acc
return test_loss, 100. * correct / total