def main(): model = VGG(depth=16, init_weights=True, cfg=None) # model = VGG_shaokai("vgg16") # model = ConvNet() # model = ResNet18() # model = torch.nn.DataParallel(model) model.load_state_dict(torch.load("./model_pruned/2019-04-09 11:14:52.016169/column-filter-fivelabels-masked_retrain/cifar10_vgg16_retrained_acc_93.960_4rhos_config_vgg16_v2.yaml.pt")) model.cuda() criterion = F.cross_entropy # criterion = CrossEntropyLossMaybeSmooth(smooth_eps=0).cuda() validate(test_loader, criterion, model) # test(model, criterion, test_loader) print("\n------------------------------\n") print('here') for name, weight in model.named_parameters(): if (len(weight.size()) == 4 and "shortcut" not in name): print(name, weight.size()) print('here now') test_column_sparsity(model) # test_chanel_sparsity(model) test_filter_sparsity(model)
def get_model(path='pretrained/ckpt.t7'): #model = nn.DataParallel(VGG('VGG16')) model = VGG('VGG16') #if True: # print(model._modules['features']._modules) #checkpoint=torch.load(path, map_location='cpu') #checkpoint=torch.load(path, map_location=lambda storage, loc:storage) #model was trained on GPU state_dict = torch.load(path)['net'] if False: print(state_dict.keys()) if False: print(list(model._modules.keys())) if use_gpu: print(model._modules['module']._modules['features']._modules) else: print(model._modules['features']._modules) if use_gpu: model = nn.DataParallel(VGG('VGG16')) else: model = VGG('VGG16') new_dict = OrderedDict() for key, val in state_dict.items(): key = key[7:] #since 'module.' has len 7 new_dict[key] = val.to('cpu') state_dict = new_dict if False: print(state_dict.keys()) model.load_state_dict(state_dict) return model
def get_vgg_net(model_folder, out_keys=['r11', 'r21', 'r31', 'r41', 'r51']): vgg_net = VGG(pool='avg', out_keys=out_keys) vgg_net.load_state_dict(torch.load(model_folder + 'vgg_conv.pth')) vgg_net.cuda() for param in vgg_net.parameters(): param.requires_grad = False return vgg_net
def test(test_model_dir, batch_size, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), model=None): transform = transforms.Compose([ transforms.Resize(32), transforms.ToTensor(), transforms.Normalize(mean, std) ]) test_dataset = torchvision.datasets.CIFAR100("../../datasets/", download=True, train=False, transform=transform) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True) if model is None: state_dict = torch.load(test_model_dir) model = VGG(state_dict['model_type'], True).cuda() model.load_state_dict(state_dict["model_state_dict"]) classes = [] accuracy = 0.0 with torch.no_grad(): tot = 0 correct = 0 for idx, (x, y) in enumerate(test_loader): x = x.cuda() y = y.cuda() out = model(x) _, predicted = torch.max(out.data, 1) for j in range(len(predicted)): tot = tot + 1 if predicted[j] == y.cpu()[j]: correct = correct + 1 print("Test Percent Finished: {}%.".format(idx * 100 / len(test_loader))) print( "{} th Picture Predicted: Correct Predict: {}. Correct Percentage: {:.2f}%" .format(tot, correct, correct * 100 / tot)) accuracy = correct / tot * 100 return accuracy
def run(): classes = ('Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral') crop_size = 44 trained_model = torch.load("C:/Users/Admin/Downloads/model_state.pth.tar") model = VGG("VGG19") model.load_state_dict(trained_model["model_weights"]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) model.eval() st.title("Facial expression recognition") img_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"]) if img_file is None: st.write('** Please upload an image **') original_image = Image.open(img_file, mode='r') st.image(original_image, use_column_width=True) model = 1 if st.button('Predict'): predict_image = detect(model, original_image) image = Image.fromarray(cv2.cvtColor(predict_image, cv2.COLOR_BGR2RGB)) st.image(image, use_column_width=True)
teacher = VGG(depth=16, init_weights=True, cfg=None) elif args.teachdepth == 19: teacher = VGG(depth=19, init_weights=True, cfg=None) else: sys.exit("vgg doesn't have those depth!") elif args.teacharch == "resnet": if args.teachdepth == 18: teacher = ResNet18() elif args.teachdepth == 50: teacher = ResNet50() else: sys.exit("resnet doesn't implement those depth!") else: sys.exit("unknown network") # teacher = ResNet18() teacher.load_state_dict(torch.load(args.teacher_path)) teacher.cuda() if args.multi_gpu: teacher = torch.nn.DataParallel(teacher) ############# criterion = CrossEntropyLossMaybeSmooth(smooth_eps=args.smooth_eps).cuda() # args.smooth = args.smooth_eps > 0.0 # args.mixup = config.alpha > 0.0 optimizer_init_lr = args.warmup_lr if args.warmup else args.lr optimizer = None if (args.optmzr == 'sgd'): optimizer = torch.optim.SGD(model.parameters(),
ap.add_argument("--output", type=str, help="Output path to save") ap.add_argument("--mode", type=str, help="mtcnn or haarcascade") args = ap.parse_args() mode = args.mode assert mode in {"mtcnn", "haarcascade"} classes = ('Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral') crop_size = 44 image_path = args.input #Load model trained_model = torch.load(args.trained_model) print("Load weight model with {} epoch".format(trained_model["epoch"])) model = VGG(args.model_name) model.load_state_dict(trained_model["model_weights"]) model.to(device) model.eval() transform_test = transforms.Compose([ transforms.TenCrop(crop_size), transforms.Lambda(lambda crops: torch.stack( [transforms.ToTensor()(crop) for crop in crops])) ]) def detect(): original_image = cv2.imread(image_path) if mode == "haarcascade": gray_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY) faces = haarcascade_detect.face_detect(gray_image)
def main(): best_acc = 0 start_epoch = args.start_epoch if not os.path.isdir(args.checkpoint): mkdir_p(args.checkpoint) trainloader = getdata(args, train=True) testloader = getdata(args, train=False) model = VGG(args.attention, args.nclass) if args.gpu: if torch.cuda.is_available(): model = model.cuda() cudnn.benchmark = True else: print( 'There is no cuda available on this machine use cpu instead.') args.gpu = False criterion = nn.CrossEntropyLoss() optimizer = '' if args.optimizer == 'sgd': optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) else: print(args.optimizer, 'is not correct') return title = 'cifar-10-' + args.attention if args.evaluate: print('\nEvaluation only') assert os.path.isfile( args.evaluate), 'Error: no checkpoint directory found!' checkpoint = torch.load(args.evaluate) model.load_state_dict(checkpoint['state_dict']) test_loss, test_acc = test(model, testloader, criterion, args.gpu) print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc)) return if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile( args.resume), 'Error: no checkpoint directory found!' args.checkpoint = os.path.dirname(args.resume) checkpoint = torch.load(args.resume) best_acc = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(args.checkpoint, state['attention'] + '-' + 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.checkpoint, state['attention'] + '-' + 'log.txt'), title=title) logger.set_names([ 'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.' ]) for epoch in range(start_epoch, args.epochs): start_time = time.time() adjust_learning_rate(optimizer, epoch) train_loss, train_acc = train(model, trainloader, criterion, optimizer, epoch, args.gpu) test_loss, test_acc = test(model, testloader, criterion, args.gpu) if sys.version[0] == '3': train_acc = train_acc.cpu().numpy().tolist()[0] test_acc = test_acc.cpu().numpy().tolist()[0] logger.append( [state['lr'], train_loss, test_loss, train_acc, test_acc]) is_best = test_acc > best_acc best_acc = max(test_acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'acc': test_acc, 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), 'attention': state['attention'], }, is_best, checkpoint=args.checkpoint) print(time.time() - start_time) print( "epoch: {:3d}, lr: {:.8f}, train-loss: {:.3f}, test-loss: {:.3f}, train-acc: {:2.3f}, test_acc:, {:2.3f}" .format(epoch, state['lr'], train_loss, test_loss, train_acc, test_acc)) logger.close() logger.plot() savefig(os.path.join(args.checkpoint, state['attention'] + '-' + 'log.eps')) print('Best acc:', best_acc)
return style_img = "./QiBashi.jpg" content_img = "./2.jpg" styleImg = load_img(style_img) contentImg = load_img(content_img) #for running on cuda styleImg = styleImg.cuda() contentImg = contentImg.cuda() vgg_directory = "./vgg_conv.pth" #path to pretrained vgg vgg_directory vgg = VGG() #print(vgg.state_dict()) vgg.load_state_dict(torch.load(vgg_directory)) for param in vgg.parameters(): param.requires_grad = False vgg.cuda() # Putting model on cuda class GramMatrix(nn.Module): def forward(self, input): b, c, h, w = input.size() f = input.view(b, c, h * w) #bxcx(hxw) # torch.bmm(batch1, batch2, out=None) # batch1 : bxmxp, batch2 : bxpxn -> bxmxn G = torch.bmm(f, f.transpose( 1, 2)) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc return G.div_(h * w)
train=False, download=True, transform=transforms.Compose([ transforms.ToTensor(), normalize, ])) testloader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=2) # Load the conv model net = VGG('VGG16') net = net.cuda() criterion_no_constrain = nn.CrossEntropyLoss() net.load_state_dict(torch.load(model_path)) net.eval() store_feat_maps(net) # store all feature maps and max pooling locations # # Build the deconv net net_decocnv = _vgg16Deconv() for idx, layer in enumerate(net.features): if isinstance(layer, nn.Conv2d): net_decocnv.features[net_decocnv.conv2deconv_indices[ idx]].weight.data = layer.weight.data if idx in net_decocnv.conv2deconv_bias_indices: net_decocnv.features[net_decocnv.conv2deconv_bias_indices[ idx]].bias.data = layer.bias.data net_decocnv = net_decocnv.cuda() net_decocnv.eval()
class counterGAN(): def __init__(self,device): self.nz = 100 self.beta1 = 0.5 self.real_label = 1 self.fake_label = 0 self.L = 100 self.device = device self.iters = 0 self.netG = Generator(self.nz).to(self.device) self.netD = Discriminator().to(self.device) self.netTarget = VGG('VGG16').to(self.device) self.netTarget.load_state_dict(torch.load('BestClassifierModel.pth',map_location=self.device)) # fixed_noise -> stores fixed generator seed for inference self.fixed_noise = torch.randn(64, self.nz, 1, 1, device=self.device) self.netG.apply(self.weights_init) self.netD.apply(self.weights_init) self.optimizerG = optim.Adam(self.netG.parameters(), lr=2e-4, betas=(self.beta1,0.999)) self.optimizerD = optim.Adam(self.netD.parameters(), lr=2e-4, betas=(self.beta1,0.999)) self.criterion = nn.BCELoss() self.criterionTarget = nn.CrossEntropyLoss() # criterionPerturbation -> norm of the generated noise self.criterionPerturbation = nn.MSELoss() def weights_init(self,m): classname = m.__class__.__name__ if classname.find('Conv') != -1: nn.init.normal_(m.weight.data, 0.0, 0.02) def save_state_dicts(self,path): torch.save({ 'netG_state_dict':self.netG.state_dict(), 'netD_state_dict':self.netD.state_dict(), 'optimizerG_state_dict':self.optimizerG.state_dict(), 'optimizerD_state_dict':self.optimizerD.state_dict() },path) def load_state_dicts(self,path): self.netG.load_state_dict(torch.load(path)['netG_state_dict']) self.netD.load_state_dict(torch.load(path)['netD_state_dict']) self.optimizerG.load_state_dict(torch.load(path)['optimizerG_state_dict']) self.optimizerD.load_state_dict(torch.load(path)['optimizerD_state_dict']) print('Loaded state dicts') def train(self,init_epoch, num_epochs,dataloaders,dataloaders_adv): G_losses = [] D_losses = [] T_losses = [] P_losses = [] Losses = [] img_list_adv = [] img_list_real = [] self.iters = 0 for epoch in range(init_epoch,num_epochs): # pbar = tqdm(total=len(dataloaders_adv['train'])) for idx, (D, D_adv) in enumerate(zip(dataloaders['train'], dataloaders_adv['train'])): # pbar.update(1) self.netG.train() self.netD.train() #Train Discriminator on adversarial image self.netD.zero_grad() #adversarial images definition image_adv = D_adv[0].to(self.device) target_labels_adv = D_adv[1].to(self.device) batch_size_adv = image_adv.size()[0] # label_adv -> target label to use for adversarial images while training discriminator label_adv = torch.full((batch_size_adv,),self.fake_label,dtype=torch.float,device=self.device) #real images definition image = D[0][:batch_size_adv,...].to(self.device) target_labels = D[1].to(self.device) batch_size = image.size()[0] # label -> target label to use for real images while training discriminator label = torch.full((batch_size,),self.real_label,dtype=torch.float,device=self.device) out_real= self.netD(image).view(-1) loss_d_real = self.criterion(out_real,label) # loss_d_real.backward() # D_x -> output of the discriminator for real images. Between (0,1) D_x = out_real.mean().item() #Train Discriminator on generated images noise = torch.randn(batch_size,self.nz,1,1,device=self.device) generated = self.netG(noise)+image_adv label_adv = torch.full((batch_size,),self.fake_label,dtype=torch.float,device=self.device) out_generated = self.netD(generated.clone().detach()).view(-1) loss_d_generated = self.criterion(out_generated,label_adv.detach()) # loss_d_generated.backward() # D_G_z1 -> output of the discriminator for generated images. Between (0,1) D_G_z1 = out_generated.mean().item() # loss_d -> total loss of discriminator loss_d = loss_d_generated+loss_d_real # self.optimizerD.step() #Train Generator self.netG.zero_grad() self.netTarget.zero_grad() label_adv = torch.full((batch_size,),self.real_label,dtype=torch.float,device=self.device) out_generated_2 = self.netD(generated.clone()).view(-1) loss_g = self.criterion(out_generated_2,label_adv) # loss_g.backward(retain_graph=True) # Calculate loss on classifier out_classifier = self.netTarget(generated.clone()) loss_c = self.criterionTarget(out_classifier,target_labels_adv) # loss_c.backward() # Calculate norm of noise generated loss_p = self.criterionPerturbation(image,generated) # loss -> final loss loss = loss_d+(10*loss_c)+loss_g-loss_p loss.backward() self.optimizerG.step() if idx%5==0 and idx!=0: # update the Discriminator every 5th step self.optimizerD.step() # self.optimizerG.step() # D_G_z2 -> output of the discriminator for generated images. Between (0,1) D_G_z2 = out_generated_2.mean().item() if idx%50==0: print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tLoss_C: %.4f\tLoss_P: %.4f\tOverall loss: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' % (epoch, num_epochs, idx, len(dataloaders['train']), loss_d.item(), loss_g.item(), loss_c.item(), loss_p.item(), loss.item(), D_x, D_G_z1, D_G_z2)) G_losses.append(loss_g.item()) D_losses.append(loss_d.item()) T_losses.append(loss_c.item()) P_losses.append(loss_p.item()) Losses.append(loss.item()) # Check how the generator is doing by saving G's output on fixed_noise if (self.iters % 50 == 0) or ((epoch == num_epochs-1) and (idx == len(dataloaders['train'])-1)): with torch.no_grad(): generated = (self.netG(self.fixed_noise)+image_adv[:64]).detach().cpu() img_list_adv.append(vutils.make_grid(generated, padding=2, normalize=True)) # TODO Visualize with torch.no_grad(): generated = (self.netG(self.fixed_noise)+image[:64]).detach().cpu() img_list_real.append(vutils.make_grid(generated, padding=2, normalize=True)) self.iters += 1 # pbar.close() if epoch%5==0 or epoch==num_epochs-1: self.visualize_images(img_list_adv,epoch,img_type='Adversarial') self.visualize_images(img_list_real,epoch,img_type='Real') self.save_state_dicts(f'BestcounterGAN_{epoch}.pth') return D_losses,G_losses, img_list_adv, img_list_real def visualize_images(self, img_list, epoch=None, img_type='Real'): if epoch==None: epoch = '' fig = plt.figure(figsize=(8,8)) plt.axis("off") plt.title("Generated Images {}".format(img_type)) plt.imshow(np.transpose(img_list[-1],(1,2,0))) plt.savefig('./images/Image_{}_{}.png'.format(epoch,img_type)) return def inference(self, images, fixed_noise=self.fixed_noise): generated = (self.netG(self.fixed_noise)+images).detach().cpu().tolist() return generated
poisoned_testloader = torch.utils.data.DataLoader(poisoned_testset, batch_size=50, shuffle=False, num_workers=2) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # Model print('==> Building model..') net = VGG('VGG11') net = net.to(device) # load model here with open("trained_checkpoint_vanilla.pt", "rb") as ckpt_file: trained_stat_dict = torch.load(ckpt_file, map_location='cuda') net.load_state_dict(trained_stat_dict) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) scheduler_multi_step = lr_scheduler.MultiStepLR(optimizer, milestones=[e for e in [151, 251]], gamma=0.1) # Training def train(epoch, trainloader): print('\nEpoch: %d' % epoch) net.train() train_loss = 0 correct = 0 total = 0 for batch_idx, (inputs, targets) in enumerate(trainloader): inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad()
[transforms.ToTensor()(crop) for crop in crops])) ]) if args.checkpoint is None: start_epoch = 0 model = VGG(args.model_name) optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) else: checkpoint = torch.load(args.checkpoint) start_epoch = checkpoint['epoch'] + 1 print('\nLoaded checkpoint from epoch %d.\n' % start_epoch) model = VGG(args.model_name) model.load_state_dict(checkpoint["model_weights"]) optimizer = checkpoint["optimizer"] if args.adjust_optim is not None: print("Adjust optimizer....") optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) data = FER2013(args.dataset_root, split="TRAIN", transform=transform_train) valid_data = FER2013(args.dataset_root, split="PUBLIC_TEST", transform=transform_test) train_loader = torch.utils.data.DataLoader(data, batch_size=args.bs, shuffle=True,
def train(self): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = VGG(self.model_type, True).to(device) loss_function = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=self.learning_rate, weight_decay=5e-4) start_epoch = 0 end_epoch = self.epochs average_loss_list = [] writer = SummaryWriter('logs') if self.load_model_dir is not None: checkpoint = torch.load(self.load_model_dir) model.load_state_dict(checkpoint['model_state_dict']) start_epoch += checkpoint['epoch'] end_epoch += checkpoint['epoch'] average_loss_list = checkpoint['average_loss_list'] for idx, loss in enumerate(average_loss_list): writer.add_scalar("Training Loss Average", loss, idx + 1) mean, std = self.compute_mean_std( datasets.CIFAR100(self.dataset_dir, train=True, download=True)) milestone = [60, 120, 160, 180, 200, 220] transform_ops = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize(mean, std) ]) train_dataset = datasets.CIFAR100(self.dataset_dir, train=True, download=True, transform=transform_ops) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_worker) train_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=milestone, gamma=0.2) model.train() running_loss = 0 running_idx = 0 running_temp_idx = 0 for epoch in range(start_epoch, end_epoch): total_loss = 0 for batch_idx, (x, y) in enumerate(train_loader): x = Variable(x.cuda()) y = Variable(y.cuda()) optimizer.zero_grad() predicted = model(x) loss = loss_function(predicted, y) loss.backward() optimizer.step() total_loss += loss.data.item() running_loss += loss.data.item() running_temp_idx += 1 if batch_idx % 100 == 0: print( 'Train Epoch: {} {:.2f}% Percent Finished. Current Loss: {:.6f}' .format(epoch + 1, 100 * batch_idx / len(train_loader), total_loss)) if running_temp_idx % 50 == 0: running_idx += 1 writer.add_scalar("Running Loss", running_loss / 100, running_idx) runninng_temp_idx = 0 running_loss = 0 writer.add_scalar("Training Loss Average", total_loss / len(train_loader), epoch + 1) print('Epoch {} Finished! Total Loss: {:.2f}'.format( epoch + 1, total_loss)) print("---------------Test Initalized!------------------") accuracy = test("", 128, model=model) writer.add_scalar("Test Accuracy", accuracy, epoch + 1) train_scheduler.step() average_loss_list.append(total_loss / len(train_loader)) if (epoch + 1) % 50 == 0: torch.save( { 'epoch': epoch + 1, 'model_state_dict': model.state_dict(), 'average_loss_list': average_loss_list, 'model_type': self.model_type }, self.model_save_dir + "vgg-checkpoint-{}.pth".format(epoch + 1)) torch.save( { 'epoch': epoch + 1, 'model_state_dict': model.state_dict(), 'average_loss_list': average_loss_list, 'model_type': self.model_type }, self.model_save_dir + "vgg.pth".format(epoch + 1))
styleImg = styleImg.unsqueeze(0) contentImg = contentImg.unsqueeze(0) styleImg,contentImg,content_iq = util.luminance_transfer(styleImg.numpy(),contentImg.numpy()) styleImg = Variable(torch.from_numpy(styleImg)) contentImg = Variable(torch.from_numpy(contentImg)) else: styleImg = load_image(opt.style_image) # 1x3x512x512 contentImg = load_image(opt.content_image) # 1x3x512x512 if(opt.cuda): styleImg = styleImg.cuda() contentImg = contentImg.cuda() ############### MODEL #################### vgg = VGG() vgg.load_state_dict(torch.load(opt.vgg_dir)) for param in vgg.parameters(): param.requires_grad = False if(opt.cuda): vgg.cuda() ########### LOSS & OPTIMIZER ########## class GramMatrix(nn.Module): def forward(self,input): b, c, h, w = input.size() f = input.view(b,c,h*w) # bxcx(hxw) # torch.bmm(batch1, batch2, out=None) # # batch1: bxmxp, batch2: bxpxn -> bxmxn # G = torch.bmm(f,f.transpose(1,2)) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc return G.div_(h*w) class styleLoss(nn.Module):
# Model print('==> Building model..') if args.net == 'vgg': net = VGG('VGG19') net = net.to(device) if device == 'cuda': net = torch.nn.DataParallel(net) cudnn.benchmark = True # Load weights from checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile(args.loadfile), 'Error: no checkpoint directory found!' checkpoint = torch.load(args.loadfile) net.load_state_dict(checkpoint['net']) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') def prune_weights(torchweights): weights = np.abs(torchweights.cpu().numpy()) weightshape = weights.shape rankedweights = weights.reshape( weights.size).argsort() #.reshape(weightshape) #print(weightshape) #print(rankedweights) max_weight = np.amax(weights) print("max weight", max_weight) num = weights.size
def train(): transform_train = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(size=32), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=True, num_workers=2) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') model = VGG(vars(args)) optimizer = torch.optim.SGD(model.parameters(), lr=args.lrate, momentum=0.9, weight_decay=5e-4) if args.use_cuda: model = model.cuda() if args.eval: model.load_state_dict(torch.load(args.model_dir)) model.eval() accuracy = model.evaluate(testloader) exit() total_size = len(trainloader) lrate = args.lrate best_score = 0.0 scores = [] for epoch in range(1, args.epochs + 1): model.train() for i, (image, label) in enumerate(trainloader): loss = model(image, label) model.zero_grad() optimizer.zero_grad() loss.backward() optimizer.step() if i % 100 == 0: print('Epoch = %d, step = %d / %d, loss = %.5f lrate = %.5f' % (epoch, i, total_size, loss, lrate)) model.eval() accuracy = model.evaluate(testloader) scores.append(accuracy) with open(args.model_dir + "_scores.pkl", "wb") as f: pkl.dump(scores, f) if best_score < accuracy: best_score = accuracy print('saving %s ...' % args.model_dir) torch.save(model.state_dict(), args.model_dir) if epoch % args.decay_period == 0: lrate *= args.decay for param_group in optimizer.param_groups: param_group['lr'] = lrate
]), } image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [train, val]} dataloders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=8, shuffle=True, num_workers=4) for x in [train, val]} dataset_sizes = {x: len(image_datasets[x]) for x in ['train', val]} use_gpu = torch.cuda.is_available() model = VGG(2) if os.path.exists(save_path): model.load_state_dict(torch.load(save_path)) if use_gpu: model = model.cuda() criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.95) scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) print ('*' * 10) print ('start training') trainval(dataloders, model, optimizer, scheduler, criterion, dataset_sizes, phase='train')
contentImg = contentImg.unsqueeze(0) styleImg, contentImg, content_iq = util.luminance_transfer( styleImg.numpy(), contentImg.numpy()) styleImg = Variable(torch.from_numpy(styleImg)) contentImg = Variable(torch.from_numpy(contentImg)) else: styleImg = load_image(opt.style_image) # 1x3x512x512 contentImg = load_image(opt.content_image) # 1x3x512x512 if (opt.cuda): styleImg = styleImg.cuda() contentImg = contentImg.cuda() ############### MODEL #################### vgg = VGG() vgg.load_state_dict(torch.load(opt.vgg_dir)) for param in vgg.parameters(): param.requires_grad = False if (opt.cuda): vgg.cuda() ########### LOSS & OPTIMIZER ########## class GramMatrix(nn.Module): def forward(self, input): b, c, h, w = input.size() f = input.view(b, c, h * w) # bxcx(hxw) # torch.bmm(batch1, batch2, out=None) # # batch1: bxmxp, batch2: bxpxn -> bxmxn # G = torch.bmm(f, f.transpose( 1, 2)) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc
class counterGAN(): def __init__(self,device): self.nz = 100 self.beta1 = 0.5 self.real_label = 1 self.fake_label = 0 self.L = 100 self.device = device self.iters = 0 self.netG = Generator(self.nz).to(self.device) self.netD = Discriminator().to(self.device) self.netTarget = VGG('VGG16').to(self.device) self.netTarget.load_state_dict(torch.load('BestClassifierModel.pth')) self.fixed_noise = torch.randn(64, self.nz, 1, 1, device=self.device) self.netG.apply(self.weights_init) self.netD.apply(self.weights_init) self.optimizerG = optim.Adam(self.netG.parameters(), lr=2e-4, betas=(self.beta1,0.999)) self.optimizerD = optim.Adam(self.netD.parameters(), lr=2e-4, betas=(self.beta1,0.999)) self.criterion = nn.BCELoss() self.criterionTarget = nn.CrossEntropyLoss() self.criterionPerturbation = nn.MSELoss() def weights_init(self,m): classname = m.__class__.__name__ if classname.find('Conv') != -1: nn.init.normal_(m.weight.data, 0.0, 0.02) def save_state_dicts(self,path): torch.save({ 'netG_state_dict':self.netG.state_dict(), 'netD_state_dict':self.netD.state_dict(), 'optimizerG_state_dict':self.optimizerG.state_dict(), 'optimizerD_state_dict':self.optimizerD.state_dict() },path) # print('Saved state dicts') def load_state_dicts(self,path): self.netG.load_state_dict(torch.load(path)['netG_state_dict']) self.netD.load_state_dict(torch.load(path)['netD_state_dict']) self.optimizerG.load_state_dict(torch.load(path)['optimizerG_state_dict']) self.optimizerD.load_state_dict(torch.load(path)['optimizerD_state_dict']) print('Loaded state dicts') def train(self,init_epoch, num_epochs,dataloader): G_losses = [] D_losses = [] T_losses = [] P_losses = [] Losses = [] img_list = [] self.iters = 0 for epoch in range(init_epoch,num_epochs): for idx,D in enumerate(dataloader['train']): self.netG.train() self.netD.train() #Train Discriminator on real image self.netD.zero_grad() image = D[0].to(self.device) target_labels = D[1].to(self.device) batch_size = image.size()[0] label = torch.full((batch_size,),self.real_label,dtype=torch.float,device=self.device) out_real = self.netD(image).view(-1) loss_d_real = self.criterion(out_real,label) # loss_d_real.backward() D_x = out_real.mean().item() #Train Discriminator on generated images noise = torch.randn(batch_size,self.nz,1,1,device=self.device) generated = self.netG(noise)+image label = torch.full((batch_size,),self.fake_label,dtype=torch.float,device=self.device) out_generated = self.netD(generated.clone().detach()).view(-1) loss_d_generated = self.criterion(out_generated,label.detach()) # loss_d_generated.backward() D_G_z1 = out_generated.mean().item() loss_d = loss_d_generated+loss_d_real # self.optimizerD.step() #Train Generator self.netG.zero_grad() self.netTarget.zero_grad() label = torch.full((batch_size,),self.real_label,dtype=torch.float,device=self.device) out_generated_2 = self.netD(generated.clone()).view(-1) loss_g = self.criterion(out_generated_2,label) # loss_g.backward(retain_graph=True) out_classifier = self.netTarget(generated.clone()) loss_c = self.criterionTarget(out_classifier,target_labels) # loss_c.backward() D_G_z2 = out_generated_2.mean().item() loss_p = self.criterionPerturbation(image,generated) loss = loss_d+loss_c+loss_g-loss_p loss.backward() self.optimizerD.step() self.optimizerG.step() # self.optimizerG.step() if idx%50==0: print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tLoss_C: %.4f\tLoss_P: %.4f\tOverall loss: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' % (epoch, num_epochs, idx, len(dataloader['train']), loss_d.item(), loss_g.item(), loss_c.item(), loss_p.item(), loss.item(), D_x, D_G_z1, D_G_z2)) G_losses.append(loss_g.item()) D_losses.append(loss_d.item()) T_losses.append(loss_c.item()) P_losses.append(loss_p.item()) Losses.append(loss.item()) # Check how the generator is doing by saving G's output on fixed_noise if (self.iters % 500 == 0) or ((epoch == num_epochs-1) and (idx == len(dataloader['train'])-1)): with torch.no_grad(): generated = (self.netG(self.fixed_noise)+image[:64]).detach().cpu() img_list.append(vutils.make_grid(generated, padding=2, normalize=True)) self.iters += 1 self.save_state_dicts('BestcounterGAN.pth') if epoch%5==0: self.visualize_images(img_list,epoch) return D_losses,G_losses,img_list def visualize_images(self,img_list,epoch=None): if epoch==None: epoch = '' fig = plt.figure(figsize=(8,8)) plt.axis("off") plt.title("Generated Images") plt.imshow(np.transpose(img_list[-1],(1,2,0))) plt.savefig('./images/Image_{}.png'.format(epoch)) return