def __init__(self, weights=False): super(VGG, self).__init__() if weights is False: self.model = models.vgg19(pretrained=True) else: self.model = models.vgg19(pretrained=False) pre = torch.load(weights) self.model.load_state_dict(pre) # output features not predictions. self.vgg19 = self.model.features for param in self.vgg19.parameters(): param.requires_grad = False
def get_vanilla_vgg_features(cut_idx=-1): if not os.path.exists('vgg_features.pth'): os.system( 'wget --no-check-certificate -N https://s3-us-west-2.amazonaws.com/jcjohns-models/vgg19-d01eb7cb.pth') vgg_weights = torch.load('vgg19-d01eb7cb.pth') # fix compatibility issues map = {'classifier.6.weight':u'classifier.7.weight', 'classifier.6.bias':u'classifier.7.bias'} vgg_weights = OrderedDict([(map[k] if k in map else k,v) for k,v in vgg_weights.iteritems()]) model = models.vgg19() model.classifier = nn.Sequential(View(), *model.classifier._modules.values()) model.load_state_dict(vgg_weights) torch.save(model.features, 'vgg_features.pth') torch.save(model.classifier, 'vgg_classifier.pth') vgg = torch.load('vgg_features.pth') if cut_idx > 36: vgg_classifier = torch.load('vgg_classifier.pth') vgg = nn.Sequential(*(vgg._modules.values() + vgg_classifier._modules.values())) vgg.eval() return vgg
def __init__(self): super(VGGNet, self).__init__() """Select conv1_1 ~ conv5_1 activation maps.""" #self.select = ['0', '5', '10', '19', '28'] self.select = ['0','2','5','7','10','12','14','16','19','21','23','25','28','30','32','34'] self.vgg = models.vgg19(pretrained=True).features
def __init__(self): super(FeatureExtractor, self).__init__() vgg19_model = vgg19(pretrained=True) # Extracts features at the 11th layer self.feature_extractor = nn.Sequential(*list(vgg19_model.features.children())[:12])
def get_pretrained_net(name): """Loads pretrained network""" if name == 'alexnet_caffe': if not os.path.exists('alexnet-torch_py3.pth'): print('Downloading AlexNet') os.system('wget -O alexnet-torch_py3.pth --no-check-certificate -nc https://box.skoltech.ru/index.php/s/77xSWvrDN0CiQtK/download') return torch.load('alexnet-torch_py3.pth') elif name == 'vgg19_caffe': if not os.path.exists('vgg19-caffe-py3.pth'): print('Downloading VGG-19') os.system('wget -O vgg19-caffe-py3.pth --no-check-certificate -nc https://box.skoltech.ru/index.php/s/HPcOFQTjXxbmp4X/download') vgg = get_vgg19_caffe() return vgg elif name == 'vgg16_caffe': if not os.path.exists('vgg16-caffe-py3.pth'): print('Downloading VGG-16') os.system('wget -O vgg16-caffe-py3.pth --no-check-certificate -nc https://box.skoltech.ru/index.php/s/TUZ62HnPKWdxyLr/download') vgg = get_vgg16_caffe() return vgg elif name == 'vgg19_pytorch_modified': # os.system('wget -O data/feature_inversion/vgg19-caffe.pth --no-check-certificate -nc https://www.dropbox.com/s/xlbdo688dy4keyk/vgg19-caffe.pth?dl=1') model = VGGModified(vgg19(pretrained=False), 0.2) model.load_state_dict(torch.load('vgg_pytorch_modified.pkl')['state_dict']) return model else: assert False
def vgg19(num_classes=1000, pretrained='imagenet'): """VGG 19-layer model (configuration "E") """ model = models.vgg19(pretrained=False) if pretrained is not None: settings = pretrained_settings['vgg19'][pretrained] model = load_pretrained(model, num_classes, settings) return model
def contentFunc(self): conv_3_3_layer = 14 cnn = models.vgg19(pretrained=True).features cnn = cnn.cuda() model = nn.Sequential() model = model.cuda() for i,layer in enumerate(list(cnn)): model.add_module(str(i),layer) if i == conv_3_3_layer: break return model
def __init__(self, descriptor_name): super(Net, self).__init__() # if descriptor_name == 'vgg16': # self.select = ['30'] # self.vgg16 = models.vgg16(pretrained=True) # self.sequence = [] # for name, layer in self.vgg16.features._modules.items(): # self.sequence += [layer] # for name, layer in self.vgg16.classifier._modules.items(): # self.sequence += [layer] # break # self.model = nn.Sequential(*self.sequence) if descriptor_name == 'vgg16': self.select = ['30'] self.vgg16 = models.vgg16(pretrained=True) self.sequence = [] for name, layer in self.vgg16.features._modules.items(): self.sequence += [layer] for name, layer in self.vgg16.classifier._modules.items(): if name == '6': break self.sequence += [layer] layer = nn.Linear(4096, 10) # init.xavier_normal(layer.weight.data, gain = 1) self.sequence += [layer] self.model = nn.Sequential(*self.sequence) elif descriptor_name == 'vgg19': self.select = ['36'] self.vgg19 = models.vgg19(pretrained=True) self.sequence = [] for name, layer in self.vgg19.features._modules.items(): self.sequence += [layer] for name, layer in self.vgg19.classifier._modules.items(): self.sequence += [layer] break self.model = nn.Sequential(*self.sequence) elif descriptor_name == 'resnet50': self.select = ['avgpool'] self.model = models.resnet50(pretrained=True) self.model.fc = nn.Linear(2048, 10) elif descriptor_name == 'resnet101': self.select = ['avgpool'] self.model = models.resnet101(pretrained=True) elif descriptor_name == 'resnet152': self.select = ['avgpool'] self.model = models.resnet152(pretrained=True) self.model.fc = nn.Linear(2048, 10)
def nst_model(content_img, style_img): vgg = models.vgg19(pretrained=True).features.eval() normalization = Normalization(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) content_img = content_img.detach() style_img = style_img.detach() content_losses = [] style_losses = [] model = nn.Sequential(normalization) i = 0 for name, layer in vgg._modules.items(): if name in ['0','2','5','10']: model.add_module('conv_{}'.format(i),layer) style_target = model(style_img) style_loss = SL(style_target) style_losses.append(style_loss) model.add_module('styleloss_{}'.format(i),style_loss) i += 1 elif name in ['7']: model.add_module('conv_{}'.format(i),layer) content_target = model(content_img) content_loss = CL(content_target) content_losses.append(content_loss) model.add_module('contentloss_{}'.format(i),content_loss) style_target = model(style_img) style_loss = SL(style_target) style_losses.append(style_loss) model.add_module('styleloss_{}'.format(i),style_loss) i += 1 elif name in ['1','3','6','8']: layer = nn.ReLU(inplace=False) model.add_module('relu_{}'.format(i),layer) i += 1 elif name in ['4','9']: model.add_module('maxpool_{}'.format(i),layer) i += 1 elif name == '11': break return model, style_losses, content_losses
def __init__(self, model): super(feature_net, self).__init__() if model == 'vgg': vgg = models.vgg19(pretrained=True) self.feature == nn.Sequential(*list(vgg.children())[:-1]) self.feature.add_module('global average', nn.AvgPool2d(9)) elif model == 'inceptionv3': inception = models.inception_v3(pretrained=True) self.feature = nn.Sequential(*list(inception.children())[:-1]) self.feature._modules.pop('13') self.feature.add_modules('global average', nn.AvgPool2d(35)) elif model == 'resnet152': resnet = models.resnet152(pretrained=True) self.feature = nn.Sequential(*list(resnet.children())[:-1])
def get_vgg19_caffe(): model = vgg19() model.classifier = nn.Sequential(View(), *model.classifier._modules.values()) vgg = model.features vgg_classifier = model.classifier names = ['conv1_1','relu1_1','conv1_2','relu1_2','pool1', 'conv2_1','relu2_1','conv2_2','relu2_2','pool2', 'conv3_1','relu3_1','conv3_2','relu3_2','conv3_3','relu3_3','conv3_4','relu3_4','pool3', 'conv4_1','relu4_1','conv4_2','relu4_2','conv4_3','relu4_3','conv4_4','relu4_4','pool4', 'conv5_1','relu5_1','conv5_2','relu5_2','conv5_3','relu5_3','conv5_4','relu5_4','pool5', 'torch_view','fc6','relu6','drop6','fc7','relu7','drop7','fc8'] model = nn.Sequential() for n, m in zip(names, list(vgg) + list(vgg_classifier)): model.add_module(n, m) model.load_state_dict(torch.load('vgg19-caffe-py3.pth')) return model
def __init__(self, requires_grad=False): super(Vgg19, self).__init__() vgg_pretrained_features = models.vgg19(pretrained=True).features self.slice1 = torch.nn.Sequential() self.slice2 = torch.nn.Sequential() self.slice3 = torch.nn.Sequential() self.slice4 = torch.nn.Sequential() self.slice5 = torch.nn.Sequential() for x in range(2): self.slice1.add_module(str(x), vgg_pretrained_features[x]) for x in range(2, 7): self.slice2.add_module(str(x), vgg_pretrained_features[x]) for x in range(7, 12): self.slice3.add_module(str(x), vgg_pretrained_features[x]) for x in range(12, 21): self.slice4.add_module(str(x), vgg_pretrained_features[x]) for x in range(21, 30): self.slice5.add_module(str(x), vgg_pretrained_features[x]) if not requires_grad: for param in self.parameters(): param.requires_grad = False
def load_model(arch='vgg19', num_labels=102, hidden_units=4096): # Load a pre-trained model if arch=='vgg19': # Load a pre-trained model model = models.vgg19(pretrained=True) elif arch=='alexnet': model = models.alexnet(pretrained=True) else: raise ValueError('Unexpected network architecture', arch) # Freeze its parameters for param in model.parameters(): param.requires_grad = False # Features, removing the last layer features = list(model.classifier.children())[:-1] # Number of filters in the bottleneck layer num_filters = model.classifier[len(features)].in_features # Extend the existing architecture with new layers features.extend([ nn.Dropout(), nn.Linear(num_filters, hidden_units), nn.ReLU(True), nn.Dropout(), nn.Linear(hidden_units, hidden_units), nn.ReLU(True), nn.Linear(hidden_units, num_labels), ##nn.Softmax(dim=1) # Please, notice Softmax layer has not been added as per Pytorch answer: # https://github.com/pytorch/vision/issues/432#issuecomment-368330817 # It is not either included in its transfer learning tutorial: # https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html ]) model.classifier = nn.Sequential(*features) return model
# print(height_s, width_s) # print(height_c, width_c) style_img = utils.image_to_tensor(style_img).unsqueeze(0) content_img = utils.image_to_tensor(content_img).unsqueeze(0) style_img = style_img.to(device, torch.float) content_img = content_img.to(device, torch.float) # print('content_img size: ', content_img.size()) # utils.show_pic(style_img, 'style image') # utils.show_pic(content_img, 'content image') # ------------------------- # Eval() means the parameters of cnn are frozen. cnn = models.vgg19(pretrained=True).features.to(config.device0).eval() cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(config.device0) cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(config.device0) # Two different initialization ways input_img = torch.randn(1, 3, height_c, width_c).to(config.device0) # input_img = content_img.clone() # print('input_img size: ', input_img.size()) output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std, content_img, style_img, input_img, style_mask_tensor, content_mask_tensor, L) print('Style transfer completed')
def reconstruct_image(img_size=224, test_case=1, num_epochs=200, print_iter=10, save_iter=50): # Load pre-trained VGG19 model to extract image features model = models.vgg19(pretrained=True) if use_gpu: model = model.cuda(gpu_id) model.eval() # Generate a random image which we will optimize if use_gpu: recon_img = Variable( 1e-1 * torch.randn(1, 3, img_size, img_size).cuda(gpu_id), requires_grad=True) else: recon_img = Variable(1e-1 * torch.randn(1, 3, img_size, img_size), requires_grad=True) # Define optimizer for previously created image optimizer = optim.SGD([recon_img], lr=1e2, momentum=0.9) # Decay learning rate by a factor of 0.1 every x epochs scheduler = lr_scheduler.StepLR(optimizer, step_size=200, gamma=0.1) # Use deep generator network to get initial image if use_dgn: print('Loading deep generator network...') # Load pre-trained model tokenizer(vocabulary) dgn = BigGAN.from_pretrained('biggan-deep-256') # Prepare an input truncation = 0.4 class_vector = np.zeros((1, 1000), dtype='float32') noise_vector = truncated_noise_sample(truncation=truncation, batch_size=1) # All in tensors class_vector = torch.from_numpy(class_vector) noise_vector = torch.from_numpy(noise_vector) if use_gpu: class_vector = class_vector.cuda(gpu_id) noise_vector = noise_vector.cuda(gpu_id) dgn.cuda(gpu_id) # Generate image with torch.no_grad(): output = dgn(noise_vector, class_vector, truncation) output = output.cpu() output = nn.functional.interpolate(output, size=(img_size, img_size), mode='bilinear', align_corners=True) if use_gpu: recon_img = Variable(output.cuda(gpu_id), requires_grad=True) else: recon_img = Variable(output, requires_grad=True) # Training for epoch in range(num_epochs): scheduler.step() optimizer.zero_grad() # Get the features from the model of the generated image output_features = get_features_from_layers(model, recon_img) # Calculate the losses euc_loss, alpha_loss, tv_loss, total_loss = get_total_loss( recon_img, output_features, original_feats) # Step total_loss.backward() optimizer.step() # Generate image every x iterations if (epoch + 1) % print_iter == 0: print('Epoch %d:\tAlpha: %.6f\tTV: %.6f\tEuc: %.6f\tLoss: %.6f' % (epoch + 1, alpha_loss.data.cpu().numpy(), tv_loss.data.cpu().numpy(), euc_loss.data.cpu().numpy(), total_loss.data.cpu().numpy())) # Save the image every x iterations if (epoch + 1) % save_iter == 0: img_sample = torch.squeeze(recon_img.cpu()) im_path = examples_dir + cls + '_' + key + '_eeg.jpg' save_image(img_sample, im_path, normalize=True)
def style_transfer(style_img, content_img, outputpath='./result.png', num_steps=500, style_weight=100000, content_weight=1, name='test', loss_dir='losses'): ''' the main function of neural style transfer :param style_img: the image with target style you want to transfer to :param content_img: the original image, to transfer its style while reserve its content :param outputpath: the path to save image with transferred style :param num_steps: number of steps to update parameters :param style_weight: weight of style :param content_weight: weight of loss ''' imsize = 512 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") transform = transforms.Compose([ transforms.Resize(imsize), transforms.CenterCrop(imsize), transforms.ToTensor() ]) style_img = image_loader(style_img, transform, device) content_img = image_loader(content_img, transform, device) # use the features module of pretrained vgg19 # need the output of the individual convolution layers to measure content and style loss. cnn = models.vgg19(pretrained=True).features.to(device).eval() # desired depth layers to compute style/content losses : content_layers_default = ['conv_4'] style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] input_img = content_img.clone() output, style_loss, content_loss = run(cnn, content_layers_default, style_layers_default, content_img, style_img, input_img, device, num_steps=num_steps, style_weight=style_weight, content_weight=content_weight) output = output.detach().cpu().numpy().squeeze(0).transpose([1, 2, 0]) plt.imsave(outputpath, output) plt.clf() x = [i * 10 for i in range(len(style_loss))] plt.plot(x, style_loss, label='style_loss') plt.plot(x, content_loss, label='content_loss') plt.xlabel('steps') plt.ylabel('loss') plt.legend() plt.savefig(os.path.join(loss_dir, "loss" + name))
def __init__(self): super(VGG19, self).__init__() vgg19 = models.vgg19(pretrained=False) # Load downloaded parameters print('Load VGG...', end=' ') vgg19.load_state_dict(torch.load('./vgg19/vgg19-dcbb9e9d.pth')) print('Load end.') features = vgg19.features self.relu1_1 = torch.nn.Sequential() self.relu1_2 = torch.nn.Sequential() self.relu2_1 = torch.nn.Sequential() self.relu2_2 = torch.nn.Sequential() self.relu3_1 = torch.nn.Sequential() self.relu3_2 = torch.nn.Sequential() self.relu3_3 = torch.nn.Sequential() self.relu3_4 = torch.nn.Sequential() self.relu4_1 = torch.nn.Sequential() self.relu4_2 = torch.nn.Sequential() self.relu4_3 = torch.nn.Sequential() self.relu4_4 = torch.nn.Sequential() self.relu5_1 = torch.nn.Sequential() self.relu5_2 = torch.nn.Sequential() self.relu5_3 = torch.nn.Sequential() self.relu5_4 = torch.nn.Sequential() for x in range(2): self.relu1_1.add_module(str(x), features[x]) for x in range(2, 4): self.relu1_2.add_module(str(x), features[x]) for x in range(4, 7): self.relu2_1.add_module(str(x), features[x]) for x in range(7, 9): self.relu2_2.add_module(str(x), features[x]) for x in range(9, 12): self.relu3_1.add_module(str(x), features[x]) for x in range(12, 14): self.relu3_2.add_module(str(x), features[x]) for x in range(14, 16): self.relu3_3.add_module(str(x), features[x]) for x in range(16, 18): self.relu3_4.add_module(str(x), features[x]) for x in range(18, 21): self.relu4_1.add_module(str(x), features[x]) for x in range(21, 23): self.relu4_2.add_module(str(x), features[x]) for x in range(23, 25): self.relu4_3.add_module(str(x), features[x]) for x in range(25, 27): self.relu4_4.add_module(str(x), features[x]) for x in range(27, 30): self.relu5_1.add_module(str(x), features[x]) for x in range(30, 32): self.relu5_2.add_module(str(x), features[x]) for x in range(32, 34): self.relu5_3.add_module(str(x), features[x]) for x in range(34, 36): self.relu5_4.add_module(str(x), features[x]) # don't need the gradients, just want the features for param in self.parameters(): param.requires_grad = False
def run_train(args): np.random.seed(args.seed) torch.manual_seed(args.seed) print('running training processing...') style_image = load_image(args.style_image, mask=False, size=args.image_style_size, scale=args.style_scale, square=True) style_image = preprocess(style_image) #save_image('style_image.png',style_image) cnn = None if args.loss_model == 'vgg19': cnn = models.vgg19(pretrained=True).features.to(device).eval() elif args.loss_model == 'vgg16': cnn = models.vgg16(pretrained=True).features.to(device).eval() # get tranform net, content losses, style losses loss_net, content_losses, style_losses, tv_loss = build_loss_model( cnn, args, style_image) #print(loss_net) #collect space back cnn = None del cnn if args.backend == 'cudnn': torch.backends.cudnn.enabled = True #this is to define the inchannels of transferm_net in_channels = 3 transform = transforms.Compose([ transforms.Resize(args.image_size), transforms.CenterCrop(args.image_size), transforms.ToTensor(), ]) train_dataset = datasets.ImageFolder(args.dataset, transform) train_loader = DataLoader(train_dataset, batch_size=args.batch_size) transform_net = TransformerNet(in_channels).to(device) mse_loss = nn.MSELoss() optimizer = optim.Adam(transform_net.parameters(), lr=args.learning_rate) iteration = [0] while iteration[0] <= args.epochs - 1: transform_net.train() agg_content_loss = 0. agg_style_loss = 0. count = 0 for batch_id, (x, _) in enumerate(train_loader): stloss = 0. ctloss = 0. n_batch = len(x) count += n_batch optimizer.zero_grad() #stack color_content_masks into x as input x, x_ori = x.to(device), x.to(device).clone() x = preprocess(x) x_ori = preprocess(x_ori) #save_image('content_x1.png', x[3].unsqueeze(0)) #assert 0 == 1 #forward input to transform_net y = transform_net(x) #compute pixel loss pixloss = 0. if args.pixel_weight > 0: pixloss = mse_loss(x, y) * args.pixel_weight #compute content loss and style loss for ctl in content_losses: ctl.mode = 'capture' for stl in style_losses: stl.mode = 'None' loss_net(x_ori) for ctl in content_losses: ctl.mode = 'loss' for stl in style_losses: stl.mode = 'loss' loss_net(y) for ctl in content_losses: ctloss += mse_loss(ctl.target, ctl.input) * args.content_weight for stl in style_losses: local_G = gram_matrix(stl.input) stloss += mse_loss(local_G, stl.target) * args.style_weight if tv_loss is not None: tvloss = tv_loss.loss else: tvloss = 0. loss = ctloss + stloss + pixloss #+ tvloss loss.backward() optimizer.step() agg_content_loss += ctloss.item() agg_style_loss += stloss.item() if (batch_id + 1) % args.log_interval == 0: mesg = "{}, Epoch {}:\t[{}/{}], content: {:.6f}, style: {:.6f}, total: {:.6f}".format( time.ctime(), iteration[0], count, len(train_dataset), agg_content_loss / (batch_id + 1), agg_style_loss / (batch_id + 1), (agg_content_loss + agg_style_loss) / (batch_id + 1)) print(mesg) if args.checkpoint_model_dir is not None and ( batch_id + 1) % args.checkpoint_interval == 0: transform_net.eval().cpu() ckpt_model_filename = "ckpt_epoch_" + str( iteration[0] + 1) + "_batch_id_" + str(batch_id + 1) + ".pth" ckpt_model_path = os.path.join(args.checkpoint_model_dir, ckpt_model_filename) torch.save(transform_net.state_dict(), ckpt_model_path) transform_net.to(device).train() iteration[0] += 1 #save final model transform_net.eval().cpu() save_model_filename = "epoch_" + str(args.epochs) + "_" + str( time.ctime()).replace(' ', '_') + "_content_" + str( args.content_weight) + "_style_" + str( args.style_weight) + ".model" save_model_path = os.path.join(args.save_model_dir, save_model_filename) torch.save(transform_net.state_dict(), save_model_path) print("\n training process is Done!, trained model saved at", save_model_path)
def main(args): #with torch.cuda.device(args.gpu): layers_map = { 'relu4_2': '22', 'relu2_2': '8', 'relu3_2': '13', 'relu1_2': '4' } vis = visdom.Visdom(port=args.display_port) loss_graph = { "g": [], "gd": [], "gf": [], "gpl": [], "gpab": [], "gs": [], "d": [], "gdl": [], "dl": [], } # for rgb the change is to feed 3 channels to D instead of just 1. and feed 3 channels to vgg. # can leave pixel separate between r and gb for now. assume user use the same weights transforms = get_transforms(args) if args.color_space == 'rgb': args.pixel_weight_ab = args.pixel_weight_rgb args.pixel_weight_l = args.pixel_weight_rgb rgbify = custom_transforms.toRGB() train_dataset = ImageFolder('train', args.data_path, transforms) train_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True) val_dataset = ImageFolder('val', args.data_path, transforms) indices = torch.randperm(len(val_dataset)) val_display_size = args.batch_size val_display_sampler = SequentialSampler(indices[:val_display_size]) val_loader = DataLoader(dataset=val_dataset, batch_size=val_display_size, sampler=val_display_sampler) # renormalize = transforms.Normalize(mean=[+0.5+0.485, +0.5+0.456, +0.5+0.406], std=[0.229, 0.224, 0.225]) feat_model = models.vgg19(pretrained=True) netG, netD, netD_local = get_models(args) criterion_gan, criterion_pixel_l, criterion_pixel_ab, criterion_style, criterion_feat, criterion_texturegan = get_criterions( args) real_label = 1 fake_label = 0 optimizerD = optim.Adam(netD.parameters(), lr=args.learning_rate_D, betas=(0.5, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=args.learning_rate, betas=(0.5, 0.999)) optimizerD_local = optim.Adam(netD_local.parameters(), lr=args.learning_rate_D_local, betas=(0.5, 0.999)) with torch.cuda.device(args.gpu): netG.cuda() netD.cuda() netD_local.cuda() feat_model.cuda() criterion_gan.cuda() criterion_pixel_l.cuda() criterion_pixel_ab.cuda() criterion_feat.cuda() criterion_texturegan.cuda() input_stack = torch.FloatTensor().cuda() target_img = torch.FloatTensor().cuda() target_texture = torch.FloatTensor().cuda() segment = torch.FloatTensor().cuda() label = torch.FloatTensor(args.batch_size).cuda() label_local = torch.FloatTensor(args.batch_size).cuda() extract_content = FeatureExtractor(feat_model.features, [layers_map[args.content_layers]]) extract_style = FeatureExtractor( feat_model.features, [layers_map[x.strip()] for x in args.style_layers.split(',')]) model = { "netG": netG, "netD": netD, "netD_local": netD_local, "criterion_gan": criterion_gan, "criterion_pixel_l": criterion_pixel_l, "criterion_pixel_ab": criterion_pixel_ab, "criterion_feat": criterion_feat, "criterion_style": criterion_style, "criterion_texturegan": criterion_texturegan, "real_label": real_label, "fake_label": fake_label, "optimizerD": optimizerD, "optimizerD_local": optimizerD_local, "optimizerG": optimizerG } for epoch in range(args.load_epoch, args.num_epoch): train(model, train_loader, val_loader, input_stack, target_img, target_texture, segment, label, label_local, extract_content, extract_style, loss_graph, vis, epoch, args)
# Importing the Model # ------------------- # # Now we need to import a pre-trained neural network. We will use a 19 # layer VGG network like the one used in the paper. # # PyTorch’s implementation of VGG is a module divided into two child # ``Sequential`` modules: ``features`` (containing convolution and pooling layers), # and ``classifier`` (containing fully connected layers). We will use the # ``features`` module because we need the output of the individual # convolution layers to measure content and style loss. Some layers have # different behavior during training than evaluation, so we must set the # network to evaluation mode using ``.eval()``. # cnn = models.vgg19(pretrained=True).features.to(device).eval() ###################################################################### # Additionally, VGG networks are trained on images with each channel # normalized by mean=[0.485, 0.456, 0.406] and std=[0.229, 0.224, 0.225]. # We will use them to normalize the image before sending it into the network. # cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device) cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device) # create a module to normalize input image so we can easily put it in a # nn.Sequential class Normalization(nn.Module):
def __init__(self): super(Vgg19, self).__init__() features = list(vgg19(pretrained=True).features) self.features = torch.nn.ModuleList(features).eval()
opts = parser.parse_args() ''' VGG19 test using examples from https://github.com/jacobgil/pytorch-grad-cam for comparison ''' from torchvision.models import vgg19 if opts.timer: import time # this block is mostly copied/adapted from the repo linked above since it's VGG preprocessing stuff x = cv2.imread(opts.imname, 1) # 1 for color image x = np.float32(cv2.resize(x, (224, 224))) / 255 # move to range [0,1] im = preprocess_image(x) im = torch.cat(opts.batchsize * [im], dim=0) # create network and set to eval mode vgg = vgg19(pretrained=True) # create GradCAM object and generae grad-CAMs GC = GradCAM(model=vgg, device=opts.device, verbose=opts.verbose) if opts.classes.__class__.__name__ == 'list': classes = opts.classes elif opts.classes.__class__.__name__ == 'int': classes = [opts.classes] else: raise Exception("Invalid list of classes provided!") # place hooks on ReLU modules inside the vgg.features submodule if opts.guided: hookmods = [] for name, module in vgg.features._modules.items(): if module.__class__.__name__ == 'ReLU':
def main(net_type='vgg'): epoch = 100 batch_size = 64 n_val = 1000 lr = 0.0001 decay = 0.00005 momentum = 0.9 cudnn.benchmark = True logger.info("Define Network and Loss...") if net_type == 'vgg': net = models.vgg19(pretrained=True) elif net_type == 'resnet': net = models.resnet50(pretrained=True) net.to(device) criterion = torch.nn.CrossEntropyLoss() opt = torch.optim.SGD(net.parameters(), lr=lr, momentum=momentum, weight_decay=decay) logger.info("Finish") traindir = '/home/lhy/ILSVRC2012/train' valdir = '/home/lhy/ILSVRC2012/val' # train_dataset = NoriInpaintDataset(train_img_nori_list_path, train_img_nori_path, mask_flist_paths_dict) # train_dataloader = train_dataset.loader(batch_size=batch_size, num_workers=batch_size//4, shuffle=True) # val_dataset = NoriInpaintDataset(val_img_nori_list_path, val_img_nori_path, mask_flist_paths_dict) # val_dataloader = val_dataset.loader(batch_size=batch_size, num_workers=batch_size//4, shuffle=False) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True) val_loader = torch.utils.data.DataLoader(datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])), batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True) for i in range(epoch): adjust_learning_rate(opt, epoch, lr) train(net, train_loader, i, opt, criterion) # evaluate on validation set prec1, perc5 = validate(net, val_loader, criterion) msg = "Epoch {}/{} Val, Perc1:{}, Perc5{}".format( i, epoch, prec1, perc5) logger.info(msg) is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'arch': net_type, 'state_dict': net.state_dict(), 'best_prec1': best_prec1, 'optimizer': opt.state_dict(), }, is_best, filename=sys.argv[1] + "checkpoint.pth.tar")
def main(): print("SSSRNet3_mask training finetuning on VOC 160*160 patches.") global opt, model, netContent opt = parser.parse_args() print(opt) gpuid = 0 cuda = True if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") opt.seed = random.randint(1, 10000) print("Random Seed: ", opt.seed) torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) cudnn.benchmark = True if opt.vgg_loss: print('===> Loading VGG model') netVGG = models.vgg19() netVGG.load_state_dict( model_zoo.load_url( 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth')) class _content_model(nn.Module): def __init__(self): super(_content_model, self).__init__() self.feature = nn.Sequential( *list(netVGG.features.children())[:-1]) def forward(self, x): out = self.feature(x) return out netContent = _content_model() print("===> Building model") model = Net() print(get_n_params(model)) model_pretrained = torch.load( 'model/model_DIV2K_noBN_96_epoch_36.pth', map_location=lambda storage, loc: storage)["model"] finetune = True if finetune == True: index = 0 for (src, dst) in zip(model_pretrained.parameters(), model.parameters()): if index > 1: list(model.parameters())[index].data = src.data index = index + 1 criterion = nn.MSELoss(size_average=False) print("===> Setting GPU") if cuda: model = model.cuda(gpuid) model_pretrained = model_pretrained.cuda(gpuid) criterion = criterion.cuda(gpuid) if opt.vgg_loss: netContent = netContent.cuda(gpuid) # optionally resume from a checkpoint if opt.resume: if os.path.isfile(opt.resume): print("=> loading checkpoint '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) opt.start_epoch = checkpoint["epoch"] + 1 model.load_state_dict(checkpoint["model"].state_dict()) else: print("=> no checkpoint found at '{}'".format(opt.resume)) # optionally copy weights from a checkpoint if opt.pretrained: if os.path.isfile(opt.pretrained): print("=> loading model '{}'".format(opt.pretrained)) weights = torch.load(opt.pretrained) model.load_state_dict(weights['model'].state_dict()) else: print("=> no model found at '{}'".format(opt.pretrained)) print("===> Setting Optimizer") optimizer = optim.Adam(model.parameters(), lr=opt.lr) print("===> Training1") #root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5' root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train_label160_HDF5' files_num = len(os.listdir(root_dir)) for epoch in range(opt.start_epoch, opt.nEpochs + 1): #save_checkpoint(model, epoch) print("===> Loading datasets") x = random.sample(os.listdir(root_dir), files_num) for index in range(0, files_num): train_path = os.path.join(root_dir, x[index]) print("===> Training datasets: '{}'".format(train_path)) train_set = DatasetFromHdf5(train_path) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) avgloss = train(training_data_loader, optimizer, model, model_pretrained, criterion, epoch, gpuid) if epoch % 2 == 0: save_checkpoint(model, epoch)
parser.add_argument('--arch', default='vgg19') parser.add_argument('--epochs', default=5) parser.add_argument('--hidden_units', default=5000) parser.add_argument('--learnrate', default=0.003) parser.add_argument('--gpu', action='store_true') args = parser.parse_args() # Load data and create data loaders train_data, valid_data, test_data = get_data(args.data_dir) trainloader, validloader, testloader = get_loader(train_data, valid_data, test_data) # Load model if args.arch == 'vgg19': model = models.vgg19(pretrained=True) if args.arch == 'vgg13': model = models.vgg13(pretrained=True) # Define hyperparameters learn_rate = float(args.learnrate) epochs = int(args.epochs) fc1_out = fc2_in = int(args.hidden_units) fc2_out = output_in = 1000 # Freeze parameters so we don't backprop through them for param in model.parameters(): param.requires_grad = False # Define new classifier to be trained on flower datset with 102 classes
def __init__(self, target_feature): super(StyleLoss, self).__init__() self.target = gram_matrix(target_feature).detach() def forward(self, input): G = gram_matrix(input) self.loss = F.mse_loss(G, self.target) return input # ----------------------- # Load the neural network # ----------------------- cnn = models.vgg19(pretrained=True).features.to(device).eval() cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device) cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device) # create a module to normalize input image so we can easily put it in a # nn.Sequential class Normalization(nn.Module): def __init__(self, mean, std): super(Normalization, self).__init__() # .view the mean and std to make them [C x 1 x 1] so that they can # directly work with image Tensor of shape [B x C x H x W]. # B is batch size. C is number of channels. H is height and W is width. self.mean = torch.tensor(mean).view(-1, 1, 1) self.std = torch.tensor(std).view(-1, 1, 1)
def __init__(self): super(Model, self).__init__() self.backbone = models.vgg19() num_features = self.backbone.classifier[6].in_features self.backbone.classifier[6] = nn.Linear(num_features, len(classes))
def initialModel(modelName, numClasses, featureExtract, usePretrained): if modelName == "vgg19": modelUse = models.vgg19(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.features._modules['0'] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) modelUse.classifier._modules['6'] = nn.Linear(in_features=4096, out_features=numClasses, bias=True) elif modelName == "shufflenet_v2_x2_0": modelUse = models.shufflenet_v2_x2_0(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.conv1._modules['0'] = nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) if dropout: modelUse.fc = nn.Sequential(nn.Dropout(dropNum), nn.Linear(2048, numClasses, bias=True)) else: modelUse.fc = nn.Linear(2048, numClasses, bias=True) elif modelName == "shufflenet_v2_x0_5": modelUse = models.shufflenet_v2_x0_5(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.conv1._modules['0'] = nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) if dropout: modelUse.fc = nn.Sequential( nn.Dropout(dropNum), nn.Linear(1024, numClasses, bias=True), ) else: modelUse.fc = nn.Linear(1024, numClasses, bias=True) elif modelName == "shufflenet_v2_x1_5": modelUse = models.shufflenet_v2_x1_5(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.conv1._modules['0'] = nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) if dropout: modelUse.fc = nn.Sequential( nn.Dropout(dropNum), nn.Linear(1024, numClasses, bias=True), ) else: modelUse.fc = nn.Linear(1024, numClasses, bias=True) elif modelName == "shufflenet_v2_x1_0": modelUse = models.shufflenet_v2_x1_0(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.conv1._modules['0'] = nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) if dropout: modelUse.fc = nn.Sequential(nn.Dropout(dropNum), nn.Linear(1024, numClasses, bias=True)) else: modelUse.fc = nn.Linear(1024, numClasses, bias=True) elif modelName == "squeezenet1_1": modelUse = models.squeezenet1_1(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.features._modules['0'] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(2, 2)) modelUse.classifier = nn.Sequential( nn.Dropout(p=dropNum, inplace=False), nn.Conv2d(512, numClasses, kernel_size=(1, 1), stride=(1, 1)), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(output_size=(1, 1))) elif modelName == "densenet121": modelUse = models.densenet121(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.features._modules['conv0'] = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) if dropout: modelUse.classifier = nn.Sequential( nn.Dropout(dropNum), nn.Linear(1024, numClasses, bias=True)) else: modelUse.classifier = nn.Linear(1024, numClasses, bias=True) elif modelName == "vgg16": modelUse = models.vgg16(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.features._modules['0'] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) modelUse.classifier._modules['6'] = nn.Linear(in_features=4096, out_features=numClasses, bias=True) elif modelName == 'mobilenet_v2': modelUse = models.mobilenet_v2(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) if mono: modelUse.features._modules['0']._modules['0'] = nn.Conv2d( 1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) modelUse.classifier = nn.Sequential( nn.Dropout(p=dropNum, inplace=False), nn.Linear(in_features=1280, out_features=numClasses, bias=True)) elif modelName == "resnet18": modelUse = models.resnet18(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) numFeatures = modelUse.fc.in_features if mono: modelUse.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) if dropout: modelUse.fc = nn.Sequential(nn.Dropout(dropNum), nn.Linear(numFeatures, numClasses)) else: modelUse.fc = nn.Linear(numFeatures, numClasses) elif modelName == "resnet34": modelUse = models.resnet34(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) numFeatures = modelUse.fc.in_features if mono: modelUse.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) if dropout: modelUse.fc = nn.Sequential(nn.Dropout(dropNum), nn.Linear(numFeatures, numClasses)) else: modelUse.fc = nn.Linear(numFeatures, numClasses) elif modelName == "resnet50": modelUse = models.resnet50(pretrained=usePretrained) setParameterRequiresGrad(modelUse, featureExtract) numFeatures = modelUse.fc.in_features if mono: modelUse.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) if dropout: modelUse.fc = nn.Sequential(nn.Dropout(dropNum), nn.Linear(numFeatures, numClasses)) else: modelUse.fc = nn.Linear(numFeatures, numClasses) else: print("model not implemented") return None, None return modelUse
def train_model(args, device, cfg=cfg): # Load VGG19 features vgg = vgg19(pretrained=True).features vgg = vgg.to(device) # We don't want to train VGG for param in vgg.parameters(): param.requires_grad_(False) # Load style net style_net = HRNet() style_net = style_net.to(device) # Load images content_img = load_image(os.path.join(args.img_root, args.content_img), size=args.content_size) content_img = content_img.to(device) style_img = load_image(os.path.join(args.img_root, args.style_img)) style_img = style_img.to(device) # Get features from VGG content_features = get_features(content_img, vgg) style_features = get_features(style_img, vgg) target = content_img.clone().requires_grad_(True).to(device) style_gram_matrixs = {layer: get_gram_matrix(style_features[layer]) for layer in style_features} optim = torch.optim.Adam(style_net.parameters(), lr=cfg['lr']) scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=cfg['step_size'], gamma=cfg['gamma']) content_loss_epoch = [] style_loss_epoch = [] total_loss_epoch = [] i = 0 output_image = content_img content_weight = args.content_weight style_weight = args.style_weight # Start training for epoch in tqdm(range(cfg['steps']+1)): scheduler.step() target = style_net(content_img).to(device) target.requires_grad_(True) target_features = get_features(target, vgg) content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2']) ** 2) style_loss = 0 for layer in style_weights: target_feature = target_features[layer] target_gram_matrix = get_gram_matrix(target_feature) style_gram_matrix = style_gram_matrixs[layer] layer_style_loss = style_weights[layer] * torch.mean((target_gram_matrix - style_gram_matrix) ** 2) b, c, h, w = target_feature.shape style_loss += layer_style_loss / (c*h*w) total_loss = content_weight * content_loss + style_weight * style_loss total_loss_epoch.append(total_loss.item()) style_loss_epoch.append(style_weight * style_loss) content_loss_epoch.append(content_weight * content_loss.item()) optim.zero_grad() total_loss.backward() optim.step() if epoch % cfg['show_every'] == 0: print("After %d criterions:" % epoch) print('Total loss: ', total_loss.item()) print('Content loss: ', content_loss.item()) print('Style loss: ', style_loss.item()) plt.imshow(im_convert(target)) plt.show() # plt.savefig(f'{args.save_dir}fig{i}.png') # i += 1 output_image = target
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") root_dir = 'dataset' # train_dir = root_dir + '/train' # valid_dir = root_dir + '/valid' batch_size = 32 # Define your transforms for the training and testing sets data_transforms = { 'train': transforms.Compose([ transforms.RandomRotation(30), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'valid': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } # Load the datasets with ImageFolder image_datasets = { x: datasets.ImageFolder(os.path.join(root_dir, x), data_transforms[x]) for x in ['train', 'valid'] } dataloaders = { x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'valid'] } dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']} class_names = image_datasets['train'].classes num_class = len(class_names) print(dataset_sizes) # # Label mapping # with open('categories.json', 'r') as f: # categories: object = json.load(f) # # Run this to test the data loader # images, labels = next(iter(dataloaders['train'])) # rand_idx = np.random.randint(len(images)) # # Print(rand_idx) # print("label: {}, class: {}, name: {}".format(labels[rand_idx].item(), # class_names[labels[rand_idx].item()], # categories[class_names[labels[rand_idx].item()]])) def build_classifier(input_features, hidden_layers, output_features): classifier = nn.conv2d() if hidden_layers is None: classifier.add_module('fc0', nn.conv2d(input_features, output_features)) else: layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:]) classifier.add_module('fc0', nn.conv2d(input_features, hidden_layers[0])) classifier.add_module('relu0', nn.ReLU()) classifier.add_module('drop0', nn.Dropout(.6)) classifier.add_module('relu1', nn.ReLU()) classifier.add_module('drop1', nn.Dropout(.5)) for i, (h1, h2) in enumerate(layer_sizes): classifier.add_module('fc' + str(i + 1), nn.Linear(h1, h2)) classifier.add_module('relu' + str(i + 1), nn.ReLU()) classifier.add_module('drop' + str(i + 1), nn.Dropout(.5)) classifier.add_module( 'output', nn.Linear(hidden_layers[-1], output_features)) return classifier def train_model(model, criterion, optimizer, sched, num_epochs): since = time.time() best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch + 1, num_epochs)) print('-' * 10) # Each epoch has a training and validation phase for phase in ['train', 'valid']: if phase == 'train': print("in training mode:") model.train() # Set model to training mode else: print("in evaluation mode:") model.eval() # Set model to evaluate mode running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) print('iterating over data') # zero the parameter gradients optimizer.zero_grad() # forward # track history if only in train print('forward propagation') with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # backward + optimize only if in training phase if phase == 'train': print('backward propagation') sched.step() loss.backward() optimizer.step() # statistics running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print('{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) # deep copy the model if phase == 'valid' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) print() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) # load best model weights model.load_state_dict(best_model_wts) return model # Create classifier model = models.vgg19(pretrained=True) for parameter in model.parameters(): parameter.requires_grad = False num_in_features = 25088 num_hidden_layers = None num_out_features = 9 new_classifier = build_classifier(num_in_features, num_hidden_layers, num_out_features) # model.classifier = new_classifier # train model criterion = nn.NLLLoss() optimizer = optim.Adam(model.classifier.parameters(), lr=0.0001) scheduler = lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1) epochs = 1 model.to(device) model = train_model(model, criterion, optimizer, scheduler, epochs) # Evaluation model.eval() accuracy = 0 for inputs, labels in dataloaders['valid']: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) # Class with the highest probability is our predicted class equality = (labels.data == outputs.max(1)[1]) # Accuracy is number of correct predictions divided by all predictions accuracy += equality.type_as(torch.FloatTensor()).mean() print("Test accuracy: {:.3f}".format(accuracy / len(dataloaders['valid']))) model.class_to_idx = image_datasets['train'].class_to_idx checkpoint = { 'input_size': dataset_sizes['train'], 'output_size': num_class, 'epochs': epochs, 'batch_size': batch_size, 'model': models.vgg19(pretrained=True), 'classifier': new_classifier, 'scheduler': scheduler, 'optimizer': optimizer.state_dict(), 'state_dict': model.state_dict(), 'class_to_idx': model.class_to_idx } torch.save(checkpoint, 'trained_model.pth')
def train_save_network(data_dir, save_dir, arch, learning_rate, hidden_units, epochs, gpu): data_dir = data_dir train_dir = data_dir + '/train' valid_dir = data_dir + '/valid' test_dir = data_dir + '/test' train_transforms = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_data = datasets.ImageFolder(train_dir, transform=train_transforms) trainloader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True) valid_transforms = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) valid_data = datasets.ImageFolder(valid_dir, transform=valid_transforms) validloader = torch.utils.data.DataLoader(valid_data, batch_size=64) if arch == 'vgg11': model = models.vgg11(pretrained=True) if arch == 'vgg13': model = models.vgg13(pretrained=True) if arch == 'vgg16': model = models.vgg16(pretrained=True) if arch == 'vgg19': model = models.vgg19(pretrained=True) for param in model.parameters(): param.requires_grad = False classifier = nn.Sequential( OrderedDict([('fc1', nn.Linear(25088, hidden_units)), ('relu', nn.ReLU()), ('fc2', nn.Linear(hidden_units, 102)), ('output', nn.LogSoftmax(dim=1))])) model.classifier = classifier criterion = nn.NLLLoss() optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate) if gpu == True: device = 'cuda' else: device = 'cpu' epochs = epochs print_every = 40 steps = 0 model.to(device) for e in range(epochs): running_loss = 0 for ii, (inputs, labels) in enumerate(trainloader): steps += 1 inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model.forward(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if steps % print_every == 0: model.eval() with torch.no_grad(): valid_loss = 0 accuracy = 0 for images, labels in validloader: images, labels = images.to(device), labels.to(device) output = model.forward(images) valid_loss += criterion(output, labels).item() ps = torch.exp(output) equality = (labels.data == ps.max(dim=1)[1]) accuracy += equality.type(torch.FloatTensor).mean() print( "Epoch: {}/{}.. ".format(e + 1, epochs), "Training Loss: {:.3f}.. ".format(running_loss / print_every), "Test Loss: {:.3f}.. ".format(valid_loss / len(validloader)), "Test Accuracy: {:.3f}".format(accuracy / len(validloader))) running_loss = 0 if save_dir == 'none': pass else: model.class_to_idx = train_data.class_to_idx checkpoint = { 'input_size': 25088, 'output_size': 102, 'hidden_layers': [hidden_units], 'state_dict': model.state_dict(), 'class_to_idx': model.class_to_idx, 'optimizer_state_dict': optimizer.state_dict } torch.save(checkpoint, save_dir)
def __init__(self): """Select conv1_1 ~ conv5_1 activation maps.""" super(VGGNet, self).__init__() self.select = ['0', '5', '10', '19', '28'] self.vgg = models.vgg19(pretrained=True).features
def main(): global opt, model, netContent opt = parser.parse_args() print(opt) cuda = opt.cuda if cuda: print("=> use gpu id: '{}'".format(opt.gpus)) os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus if not torch.cuda.is_available(): raise Exception( "No GPU found or Wrong gpu id, please run without --cuda") opt.seed = random.randint(1, 10000) print("Random Seed: ", opt.seed) torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) cudnn.benchmark = True print("===> Loading datasets") training_data_loaders = [] random_order = [i for i in range(4)] for i in range(4): filename = "../train/DIV2K_train_320_HDF5/DIV2K_x4_Part" + str( i + 1) + ".h5" train_set = DatasetFromHdf5(filename) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, \ batch_size=opt.batchSize, shuffle=True) training_data_loaders.append(training_data_loader) if opt.vgg_loss: print('===> Loading VGG model') netVGG = models.vgg19() netVGG.load_state_dict( model_zoo.load_url( 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth')) class _content_model(nn.Module): def __init__(self): super(_content_model, self).__init__() self.feature = nn.Sequential( *list(netVGG.features.children())[:-1]) def forward(self, x): out = self.feature(x) return out netContent = _content_model() print("===> Building model") model = _NetG() criterion = nn.MSELoss(size_average=False) print("===> Setting GPU") if cuda: model = model.cuda() criterion = criterion.cuda() if opt.vgg_loss: netContent = netContent.cuda() # optionally resume from a checkpoint if opt.resume: if os.path.isfile(opt.resume): print("=> loading checkpoint '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) opt.start_epoch = checkpoint["epoch"] + 1 model.load_state_dict(checkpoint["model"].state_dict()) else: print("=> no checkpoint found at '{}'".format(opt.resume)) # optionally copy weights from a checkpoint if opt.pretrained: if os.path.isfile(opt.pretrained): print("=> loading model '{}'".format(opt.pretrained)) weights = torch.load(opt.pretrained) model.load_state_dict(weights['model'].state_dict()) else: print("=> no model found at '{}'".format(opt.pretrained)) print("===> Setting Optimizer") optimizer = optim.Adam(model.parameters(), lr=opt.lr) print("===> Training") for epoch in range(opt.start_epoch, opt.nEpochs + 1): shuffle(random_order) for idx in random_order: train(training_data_loaders[idx], optimizer, model, criterion, epoch) save_checkpoint(model, epoch)
import torch import torch.nn as nn import torchvision.models as models import loss vgg = models.vgg19(pretrained=True).features if torch.cuda.is_available(): vgg = vgg.cuda() content_layers_default = ['conv_4'] style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] def get_style_model_and_loss(style_img, content_img, cnn=vgg, style_weight=1000, content_weight=1, content_layers=content_layers_default, style_layers=style_layers_default): content_loss_list = [] style_loss_list = [] model = nn.Sequential() if torch.cuda.is_available(): model = model.cuda() gram = loss.Gram() if torch.cuda.is_available(): gram = gram.cuda()
def get_args(): parser = argparse.ArgumentParser() parser.add_argument('--use-cuda', action='store_true', default=False) parser.add_argument('--image-path', type=str, default='./examples/000004.jpg') args = parser.parse_args() args.use_cuda = args.use_cuda and torch.cuda.is_available() return args if __name__ == '__main__': args = get_args() grad_cam = GradCam(model=models.vgg19(pretrained=True), \ target_layer_names=["35"], use_cuda=args.use_cuda) img_name = args.image_path.split('/')[-1].split('.')[0] img = cv2.imread(args.image_path, 1) img = np.float32(cv2.resize(img, (224, 224))) / 255 input = preprocess_image(img) # Find tok and plot noutput = grad_cam.forward(input).cpu().data.numpy()[0] topk_idxs = noutput.argsort()[-5:][::-1] print(noutput[topk_idxs]) # If None, returns the map for the highest scoring category. # Otherwise, targets the requested index. # target_index = None for topl, target_index in enumerate(topk_idxs): print(f'{topl}:{imgnet_classes[target_index]}')
def __init__(self): super(MeanActivations, self).__init__() self.vgg = models.vgg19(pretrained=True).features
def __init__(self): self.cnn = models.vgg19(pretrained=True).features.to(device).eval() self.num_steps = 70 self.image_size = 100 self.progress_lambda = None self.should_terminate_lambda = None
# Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. #---------------------------------------------------------------------------------------------- import argparse import os from six import text_type as _text_type from mmdnn.conversion.examples.imagenet_test import TestKit import torch import torchvision.models as models NETWORKS_MAP = { 'inception_v3' : lambda : models.inception_v3(pretrained=True), 'vgg16' : lambda : models.vgg16(pretrained=True), 'vgg19' : lambda : models.vgg19(pretrained=True), 'resnet152' : lambda : models.resnet152(pretrained=True), 'densenet' : lambda : models.densenet201(pretrained=True), 'squeezenet' : lambda : models.squeezenet1_1(pretrained=True) } def _main(): parser = argparse.ArgumentParser() parser.add_argument('-n', '--network', type=_text_type, help='Model Type', required=True, choices=NETWORKS_MAP.keys()) parser.add_argument('-i', '--image', type=_text_type, help='Test Image Path')
def main(): # parse options # parser = TrainShapeMatchingOptions() # opts = parser.parse() # create model print('--- create model ---') netShapeM = ShapeMatchingGAN(opts.GS_nlayers, opts.DS_nlayers, opts.GS_nf, opts.DS_nf, opts.GT_nlayers, opts.DT_nlayers, opts.GT_nf, opts.DT_nf, opts.gpu) if opts.gpu: netShapeM.cuda() netShapeM.init_networks(weights_init) netShapeM.train() if opts.style_loss: netShapeM.G_S.load_state_dict(torch.load(opts.load_GS_name)) netShapeM.G_S.eval() VGGNet = models.vgg19(pretrained=True).features VGGfeatures = VGGFeature(VGGNet, opts.gpu) for param in VGGfeatures.parameters(): param.requires_grad = False if opts.gpu: VGGfeatures.cuda() style_targets = get_GRAM(opts.style_name, VGGfeatures, opts.batchsize, opts.gpu) print('--- training ---') # load image pair _, X, Y, Noise = load_style_image_pair(opts.style_name, gpu=opts.gpu) Y = to_var(Y) if opts.gpu else Y X = to_var(X) if opts.gpu else X Noise = to_var(Noise) if opts.gpu else Noise for epoch in range(opts.texture_step1_epochs): for i in range(opts.Ttraining_num // opts.batchsize): x, y = cropping_training_batches(X, Y, Noise, opts.batchsize, opts.Tanglejitter, opts.subimg_size, opts.subimg_size) losses = netShapeM.texture_one_pass(x, y) print('Step1, Epoch [%02d/%02d][%03d/%03d]' % (epoch + 1, opts.texture_step1_epochs, i + 1, opts.Ttraining_num // opts.batchsize), end=': ') print('LDadv: %+.3f, LGadv: %+.3f, Lrec: %+.3f, Lsty: %+.3f' % (losses[0], losses[1], losses[2], losses[3])) if opts.style_loss: fnames, _ = custom_load_train_batchfnames(opts.text_path, opts.augment_text_path, opts.batchsize, opts.text_datasize, opts.augment_text_datasize, trainnum=opts.Ttraining_num) for epoch in range(opts.texture_step2_epochs): itr = 0 for fname in fnames: itr += 1 t = prepare_text_batch(fname, anglejitter=False) x, y = cropping_training_batches(X, Y, Noise, opts.batchsize, opts.Tanglejitter, opts.subimg_size, opts.subimg_size) t = to_var(t) if opts.gpu else t losses = netShapeM.texture_one_pass(x, y, t, 0, VGGfeatures, style_targets) print('Step2, Epoch [%02d/%02d][%03d/%03d]' % (epoch + 1, opts.texture_step2_epochs, itr, len(fnames)), end=': ') print('LDadv: %+.3f, LGadv: %+.3f, Lrec: %+.3f, Lsty: %+.3f' % (losses[0], losses[1], losses[2], losses[3])) print('--- save ---') # directory netShapeM.save_texture_model(opts.save_path, opts.save_name)
def __init__(self): super(VGG19, self).__init__() features = models.vgg19(pretrained=True).features self.relu1_1 = torch.nn.Sequential() self.relu1_2 = torch.nn.Sequential() self.relu2_1 = torch.nn.Sequential() self.relu2_2 = torch.nn.Sequential() self.relu3_1 = torch.nn.Sequential() self.relu3_2 = torch.nn.Sequential() self.relu3_3 = torch.nn.Sequential() self.relu3_4 = torch.nn.Sequential() self.relu4_1 = torch.nn.Sequential() self.relu4_2 = torch.nn.Sequential() self.relu4_3 = torch.nn.Sequential() self.relu4_4 = torch.nn.Sequential() self.relu5_1 = torch.nn.Sequential() self.relu5_2 = torch.nn.Sequential() self.relu5_3 = torch.nn.Sequential() self.relu5_4 = torch.nn.Sequential() for x in range(2): self.relu1_1.add_module(str(x), features[x]) for x in range(2, 4): self.relu1_2.add_module(str(x), features[x]) for x in range(4, 7): self.relu2_1.add_module(str(x), features[x]) for x in range(7, 9): self.relu2_2.add_module(str(x), features[x]) for x in range(9, 12): self.relu3_1.add_module(str(x), features[x]) for x in range(12, 14): self.relu3_2.add_module(str(x), features[x]) for x in range(14, 16): self.relu3_3.add_module(str(x), features[x]) for x in range(16, 18): self.relu3_4.add_module(str(x), features[x]) for x in range(18, 21): self.relu4_1.add_module(str(x), features[x]) for x in range(21, 23): self.relu4_2.add_module(str(x), features[x]) for x in range(23, 25): self.relu4_3.add_module(str(x), features[x]) for x in range(25, 27): self.relu4_4.add_module(str(x), features[x]) for x in range(27, 30): self.relu5_1.add_module(str(x), features[x]) for x in range(30, 32): self.relu5_2.add_module(str(x), features[x]) for x in range(32, 34): self.relu5_3.add_module(str(x), features[x]) for x in range(34, 36): self.relu5_4.add_module(str(x), features[x]) # don't need the gradients, just want the features for param in self.parameters(): param.requires_grad = False
import sys from coco_dataset import * import numpy as np import torch import torch.nn as nn from torch.autograd import Variable import torchvision.models as models model = models.vgg19(pretrained=True) new_classifier = nn.Sequential(*list(model.classifier.children())[:-1]) model.classifier = new_classifier model.cuda() mytrans = transforms.Compose([Rescale((224,224)),ToVGGTensor()]) train_dir = '/home/datasets/coco/raw/train2014/' val_dir = '/home/datasets/coco/raw/val2014' #dataset = COCOImageDataset(train_dir,val_dir,transform=default_transform(224)) dataset = COCOImageCropDataset(train_dir,val_dir,transform=croped_transform()) dataloader = data.DataLoader(dataset, batch_size=64, shuffle=False, num_workers=4) for i, batch in enumerate(dataloader): fns,imgs = batch imgs = imgs.cuda() fns = list(fns) feats = [] for j,fn in enumerate(fns): feat = model.forward(Variable(imgs[j,:])) feat = torch.mean(feat.data,dim=0) feats.append(feat)
parser.add_argument('--style_weight', '-s_w', type=int, default=500, help='The weight of style loss') parser.add_argument('--initialize_noise', '-i_n', action='store_true', help='Initialize with white noise? elif initialize with content image') parser.add_argument('--cuda', action='store_true',default=True, help='use cuda?') args = parser.parse_args() use_cuda = torch.cuda.is_available() and args.cuda dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor # desired size of the output image imsize = 512 if use_cuda else 128 # use small size if no gpu #assert style_img.size() == content_img.size(), \ # "we need to import style and content images of the same size" cnn = models.vgg19(pretrained=True).features # move it to the GPU if possible: if use_cuda: cnn = cnn.cuda() # desired depth layers to compute style/content losses : content_layers_default = ['conv_4'] style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] def get_style_model_and_losses(cnn, style_img, content_img, style_weight=1000, content_weight=1, content_layers=content_layers_default, style_layers=style_layers_default): cnn = copy.deepcopy(cnn) # just in order to have an iterable access to or list of content/syle # losses
def train_model(dat_dir, arch, hid_unit, lr, epoch, device, save_name): #check whether inputs are valid if os.path.isdir(dat_dir) == False: print('Please enter a valid path to the pictures') return if (hid_unit < 102) | (hid_unit > 25088): print( 'Please enter a valid number of Hidden Units: 102 < Hid_Units < 25088' ) return if lr <= 0: print('Learning Rate must be >0') return if epoch < 0: print('Epochs number must be >0') return if (device != 'cpu') & (device != 'cuda'): print('Please choose "cpu"/"cuda" for --gpu') return # Imports import numpy as np import torch from torch import nn from torch import optim import torch.nn.functional as F from torchvision import datasets, transforms, models data_dir = dat_dir train_dir = data_dir + '/train' valid_dir = data_dir + '/valid' test_dir = data_dir + '/test' # TODO: Define your transforms for the training, validation, and testing sets train_transforms = transforms.Compose([ transforms.RandomRotation(30), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) validation_transforms = transforms.Compose([ transforms.Resize(255), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) test_transforms = transforms.Compose([ transforms.Resize(255), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # TODO: Load the datasets with ImageFolder train_data = datasets.ImageFolder(train_dir, transform=train_transforms) validation_data = datasets.ImageFolder(valid_dir, transform=train_transforms) test_data = datasets.ImageFolder(test_dir, transform=test_transforms) # TODO: Using the image datasets and the trainforms, define the dataloaders trainloader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True) validationloader = torch.utils.data.DataLoader(validation_data, batch_size=64, shuffle=False) testloader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=False) if arch == 'vgg13': model = models.vgg13(pretrained=True) elif arch == 'vgg16': model = models.vgg16(pretrained=True) elif arch == 'vgg19': model = models.vgg19(pretrained=True) else: print('Please choose VGG13/VGG16/VGG19 model') return # Freeze model parameters and weights for param in model.parameters(): param.requires_grad = False # Define 3 output layers from collections import OrderedDict classifier = nn.Sequential( OrderedDict([('fc1', nn.Linear(25088, hid_unit)), ('relu1', nn.ReLU()), ('dropout1', nn.Dropout(p=0.2)), ('fc4', nn.Linear(hid_unit, 102)), ('output', nn.LogSoftmax(dim=1))])) # Assign new output layers to a model model.classifier = classifier #Define Loss Criterion and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.classifier.parameters(), lr=lr) # Send Model to Device model.to(device) # Loop over Epochs for e in range(epoch): # Null the Stat variables running_loss = 0 test_loss = 0 accuracy = 0 # Loop over Training Batches for images, labels in trainloader: # Switch on Train Mode to use gradients model.train() #Send images and labels to GPU/CPU images, labels = images.to(device), labels.to(device) # NN Routine optimizer.zero_grad() output = model.forward(images) loss = criterion(output, labels) loss.backward() optimizer.step() # Calculate Stat running_loss += loss.item() # Validation Step else: #Switch off the Gradient calculations (to increase speed) with torch.no_grad(): # Evaluation Regime (to increase speed) model.eval() #Loop over test Batches for images1, labels1 in validationloader: #Send images and labels to GPU / CPU images1, labels1 = images1.to(device), labels1.to(device) #Forward pass through the model and get the probabilities ps = torch.exp(model(images1)) #Get classes and probabilities from ts top_p, top_class = ps.topk(1, dim=1) #Calculate Loss loss = criterion(model(images1), labels1) # Statistics test_loss += loss.item() equals = top_class == labels1.view(*top_class.shape) accuracy += torch.mean(equals.type(torch.FloatTensor)) print(f'Epoch No: {e+1}') print( f'Accuracy Validation: {accuracy.item()/len(validationloader)*100}%' ) print(f'Training Loss: {running_loss/len(trainloader)}') print(f'Validation Loss: {test_loss/len(validationloader)}') #Save the model model.class_to_idx = train_data.class_to_idx torch.save( { 'architecture': arch, 'classifier': model.classifier, 'state_dict': model.state_dict(), 'class_to_idx': model.class_to_idx }, save_name)
def test_vgg19(self): process_model(models.vgg19(self.pretrained), self.image, _C_tests.forward_vgg19, "VGG19")
def __init__(self): super(VGGNet, self).__init__() self.vgg = models.vgg19(pretrained=True).features
opt = parser.parse_args() EPOCH = opt.epochs DIM = opt.input_dimension BATCH_SIZE = opt.batch_size LR = opt.learning_rate cuda_available = torch.cuda.is_available() device = torch.device("cuda" if cuda_available else "cpu") print('Using device:%s' % (device)) content_layers = ['conv_5'] # HAO['conv_4'] style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] vgg = vgg19(pretrained=True).to(device) vgg.eval() for p in vgg.parameters(): p.requires_grad = False def checkpoint_path(epoch, batch): if batch > 0: return '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/worker_e_%d_b_%d' % ( epoch, batch ), '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/judge_e_%d_b_%d' % ( epoch, batch) else: return '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/worker_e_%d' % ( epoch ), '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/judge_e_%d' % (
def preprocess(self, img_path): temp_img = Image.open(img_path) return self.transform(temp_img) def __getitem__(self, index): sr_img = self.preprocess( os.path.join(self.sr_path, self.sr_names[index])) gt_img = self.preprocess( os.path.join(self.gt_path, self.gt_names[index])) return sr_img.to(device), gt_img.to(device) def __len__(self): return len(self.sr_names) vgg19 = models.vgg19(pretrained=True).to(device) vgg19.eval() vgg_data = VggDataset(sr_path='../results/RRDB_ESRGAN_x4/div-2k-test', gt_path='../data_samples/div_train/valid_HR') vgg_loader = DataLoader(vgg_data, batch_size=2) for sr, gt in vgg_loader: sr_vgg = vgg19(sr) gt_vgg = vgg19(gt) vgg_loss = ((gt_vgg - sr_vgg)**2).mean() print(vgg_loss)
def __init__(self, config=GlobalConfig()): super(HighQualityNet, self).__init__() self.config = config """ Settings """ if self.config.architecture == 'AlexNet': self.classification_net = models.alexnet(pretrained=True).cuda() print(self.classification_net) elif self.config.architecture == 'ResNet': self.classification_net = models.resnet50(pretrained=True).cuda() print(self.classification_net) elif self.config.architecture == 'VGG': self.classification_net = models.vgg19(pretrained=True).cuda() print(self.classification_net) elif self.config.architecture == 'DenseNet': self.classification_net = models.densenet121( pretrained=True).cuda() print(self.classification_net) elif self.config.architecture == 'ResNet': self.classification_net = models.resnet152(pretrained=True).cuda() print(self.classification_net) elif self.config.architecture == 'GoogleNet': self.classification_net = models.googlenet(pretrained=True).cuda() print(self.classification_net) else: self.classification_net = models.mobilenet_v2( pretrained=True).cuda() print(self.classification_net) if torch.cuda.device_count() > 1: self.classification_net = torch.nn.DataParallel( self.classification_net) self.criterion = nn.CrossEntropyLoss().cuda() self.encoder = Prep_pureUnet(config=config).cuda() if torch.cuda.device_count() > 1: self.encoder = torch.nn.DataParallel(self.encoder) print(self.encoder) self.optimizer = torch.optim.Adam(self.encoder.parameters()) """ Noise Layers """ self.noise_layers = [Identity()] # self.cropout_layer = Cropout(config).cuda() self.jpeg_layer_80 = DiffJPEG(256, 256, quality=80, differentiable=True).cuda() self.jpeg_layer_90 = DiffJPEG(256, 256, quality=90, differentiable=True).cuda() self.jpeg_layer_70 = DiffJPEG(256, 256, quality=70, differentiable=True).cuda() self.jpeg_layer_60 = DiffJPEG(256, 256, quality=60, differentiable=True).cuda() self.jpeg_layer_50 = DiffJPEG(256, 256, quality=50, differentiable=True).cuda() # self.gaussian = Gaussian().cuda() # self.dropout = Dropout(self.config,keep_ratio_range=(0.5,0.75)).cuda() # self.resize = Resize().cuda() # self.crop_layer = Crop((0.2, 0.5), (0.2, 0.5)).cuda() self.noise_layers.append(self.jpeg_layer_80) self.noise_layers.append(self.jpeg_layer_90) self.noise_layers.append(self.jpeg_layer_70) self.noise_layers.append(self.jpeg_layer_60) self.noise_layers.append(self.jpeg_layer_50)