def main(): global args args = parser.parse_args() model_selection = 'resnet' model = define_model(encoder=model_selection) original_model2 = net_mask.drn_d_22(pretrained=True) model2 = net_mask.AutoED(original_model2) model = torch.nn.DataParallel(model).cuda() model2 = torch.nn.DataParallel(model2).cuda() model.load_state_dict( torch.load('./pretrained_model/model_' + model_selection)) model2.load_state_dict(torch.load('./net_mask/mask_' + model_selection)) test_loader = loaddata.getTestingData(1) test(test_loader, model, model2, 'mask_' + model_selection)
def main(): Encoder = modules.E_resnet(resnet.resnet50(pretrained=True)) N = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) G_adv = net_mask.G(net_mask.drn_d_22(pretrained=True)) N = torch.nn.DataParallel(N).cuda() G_adv = torch.nn.DataParallel(G_adv).cuda() N.load_state_dict(torch.load('./models/N')) G_adv.load_state_dict(torch.load('./models/G_adv')) cudnn.benchmark = True test_loader = loaddata.getTestingData(8) #test for N(x*G_adv(x*)) test_G_adv(test_loader, N, G_adv, epsilon=0.05, iteration=10) test_G_adv(test_loader, N, G_adv, epsilon=0.1, iteration=10) test_G_adv(test_loader, N, G_adv, epsilon=0.15, iteration=10) test_G_adv(test_loader, N, G_adv, epsilon=0.2, iteration=10)
def main(): global args args = parser.parse_args() model_selection = 'resnet' model = define_model(encoder=model_selection) original_model2 = net_mask.drn_d_22(pretrained=True) model2 = net_mask.AutoED(original_model2) if torch.cuda.device_count() == 8: model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3, 4, 5, 6, 7]).cuda() model2 = torch.nn.DataParallel(model2, device_ids=[0, 1, 2, 3, 4, 5, 6, 7]).cuda() batch_size = 64 elif torch.cuda.device_count() == 4: model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3]).cuda() model2 = torch.nn.DataParallel(model2, device_ids=[0, 1, 2, 3]).cuda() batch_size = 32 else: model = torch.nn.DataParallel(model).cuda() model2 = torch.nn.DataParallel(model2).cuda() batch_size = 8 model.load_state_dict( torch.load('./pretrained_model/model_' + model_selection)) cudnn.benchmark = True optimizer = torch.optim.Adam(model2.parameters(), args.lr, weight_decay=args.weight_decay) train_loader = loaddata.getTrainingData(batch_size) for epoch in range(args.start_epoch, args.epochs): train(train_loader, model, model2, optimizer, epoch) torch.save(model.state_dict(), '/net_mask/mask_' + model_selection)
original_model = senet.senet154(pretrained='imagenet') Encoder = modules.E_senet(original_model) model = net.model(Encoder, num_features=2048, block_channel = [256, 512, 1024, 2048]) return model def convert_to_onnx(net, output_name): input = torch.randn(1, 3, 224, 224) input_names = ['data'] output_names = ['output'] net.eval() torch.onnx.export(net, input, output_name, verbose=True, input_names=input_names, output_names=output_names,opset_version=9) model_selection = 'senet' model = define_model(encoder = model_selection) original_model2 = net_mask.drn_d_22(pretrained=True) model2 = net_mask.AutoED(original_model2) #model = model.cuda() #model2 = model2.cuda() #model = torch.nn.DataParallel(model).cuda() #model2 = torch.nn.DataParallel(model2).cuda() dic = torch.load('./pretrained_model/model_' + model_selection) print(dic.keys()) D2 = {} for k in dic.keys(): D2[k.replace("module.","",1)] = dic[k] print(D2.keys()) model.load_state_dict(D2) convert_to_onnx(model,"model.onnx")