def main(args):
# Build Models
encoder = CNN(args.hidden_size)
encoder.eval() # evaluation mode (BN uses moving mean/variance)
decoder = LSTM(args.embed_size, args.hidden_size,
len(vocab), args.num_layers)
# Load the trained model parameters
encoder.load_state_dict(torch.load(args.encoder_path))
decoder.load_state_dict(torch.load(args.decoder_path))
# load data set
is_training = True
testing_data = IDDataset(not is_training)
# If use gpu
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
test_acc = evaluation(testing_data, encoder, decoder)
print("Accuracy is %.4f" % test_acc)
def main():
# parser是训练和测试的一些参数设置，如果default里面有数值，则默认用它，
# 要修改可以修改default，也可以在命令行输入
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--model', default='CNN',#这里选择你要训练的模型
help='CNN or MLP')
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
help='SGD momentum (default: 0.5)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=50, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=True,
help='For Saving the current Model')
parser.add_argument('--save_dir', default='output/',#模型保存路径
help='dir saved models')
args = parser.parse_args()
#torch.cuda.is_available()会判断电脑是否有可用的GPU,没有则用cpu训练
use_cuda = not args.no_cuda and torch.cuda.is_available()
print(use_cuda)
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.FashionMNIST('./fashionmnist_data/', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.FashionMNIST('./fashionmnist_data/', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=args.test_batch_size, shuffle=True, **kwargs)
writer=SummaryWriter()#用于记录训练和测试的信息:loss,acc等
if args.model=='CNN':
model = CNN().to(device)#CNN() or MLP
if args.model=='MLP':
model = MLP().to(device)#CNN() or MLP
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) #optimizer存储了所有parameters的引用，每个parameter都包含gradient
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[12, 24], gamma=0.1) #学习率按区间更新
model.train()
log_loss=0
log_acc=0
for epoch in range(1, args.epochs + 1):
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target) # negative log likelihood loss(nll_loss), sum up batch cross entropy
loss.backward()
optimizer.step() # 根据parameter的梯度更新parameter的值
# 这里设置每args.log_interval个间隔打印一次训练信息，同时进行一次验证，并且将验证(测试)的准确率存入writer
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
#下面是模型验证过程
model.eval()
test_loss = 0
correct = 0
with torch.no_grad(): # 无需计算梯度
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
writer.add_scalars('loss', {'train_loss':loss,'val_loss':test_loss},global_step=log_acc)
writer.add_scalar('val_accuracy', correct / len(test_loader.dataset), global_step=log_acc)
log_acc += 1
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
model.train()
if (args.save_model):#保存训练好的模型
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
torch.save(model.state_dict(), os.path.join(args.save_dir,args.model+".pt"))
writer.add_graph(model, (data,))# 将模型结构保存成图，跟踪数据流动
writer.close()
import torch
from net import MLP, CNN #
from torchvision import datasets, transforms
from sklearn.metrics import multilabel_confusion_matrix
#
test_loader = torch.utils.data.DataLoader(datasets.FashionMNIST(
'./fashionmnist_data/',
train=False,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])),
batch_size=1,
shuffle=True)
model = CNN()
device = torch.device('cpu')
model = model.to(device)
model.load_state_dict(torch.load('output/CNN.pt'))
model.eval()
pres = []
labels = []
i = 0
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
pred = output.argmax(
dim=1, keepdim=True) # get the index of the max log-probability
pres.append(pred[0][0].item())
labels.append(target[0].item())
mcm = multilabel_confusion_matrix(labels, pres) #mcm
print(mcm)
x=x.permute(1, 2, 0)
cv2.imwrite(os.path.join(f_dir, w_dir), 255 * x.mean(2).unsqueeze(2).cpu().detach().numpy())
def feature_visualization(net,x):
w_dirs=['layer_'+str(i)+'.jpg' for i in range(1,5)]
x=net.conv1(x)
plot_x(x,w_dirs[0])
x = net.conv2(x)
plot_x(x, w_dirs[1])
x = net.conv3(x)
plot_x(x, w_dirs[2])
x = net.conv4(x)
plot_x(x, w_dirs[3])
model=CNN()#
device=torch.device('cpu')#用cpu进行推理
model=model.to(device)
model.load_state_dict(torch.load('output/CNN.pt'))##
model.eval()#告诉模型验证
print(model)#可以打印网络结构观察
#--------以上就是推理之前模型的导入--------
data_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
img=Image.open('test_image.jpg')#用于推理的图片
image=data_transforms(img)#预处理，转成tensor同时正则化
image=image.unsqueeze(0)#[1,28,28]->[1,1,28,28]
feature_visualization(model,image.to(device))#可视化
img = orig.copy().astype(np.float32)
perturbation = np.empty_like(orig)
mean = [0.5]
std = [0.5]
img /= 255.0
img = (img - mean) / std
# load model
model1 = CNN(1, 10)
saved1 = torch.load('relu.pkl', map_location='cpu')
model1.load_state_dict(saved1)
model1.eval()
criterion = nn.CrossEntropyLoss()
device = 'cuda' if gpu else 'cpu'
# prediction before attack
inp = Variable(
torch.from_numpy(img).to(device).float().unsqueeze(0).unsqueeze(0),
requires_grad=True)
out1 = model1(inp)
pred1 = np.argmax(out1.data.cpu().numpy())
print('Prediction before attack: %s' % (pred1))