def augmentation_pipeline(image, mask): alt_background = TF.to_tensor(Image.open("unicorn.jpg")) # doing random crop/pad before anything else means that background # replacement will never be cropped/padded. this prevents any 0 # padding appearing when the image is downscaled then padded back # to 218x178 crop_aug = K.RandomCrop((218,178), pad_if_needed=True) rand_scale = torch.rand(1).item()*.4 + .8 # scale between .8x - 1.2x new_dimensions = (int(218*rand_scale),int(178*rand_scale)) # resize first augmented = kornia.resize(image.unsqueeze(0), new_dimensions) aug_mask = kornia.resize(mask.unsqueeze(0), new_dimensions) # crop back to 218x178 # for some reason the generate_parameters() function for K.RandomCrop # always generates the same crop box so I have to concatenate the # image and mask so I can crop them both with the same random params aug_mask = aug_mask.type(torch.float32) concat = crop_aug(torch.cat((augmented,aug_mask), axis=1)).squeeze(0) augmented = concat[:3] aug_mask = concat[3:] # augmentation pipeline: #augmented = background_removal(image_tensor, mask) #augmented = eye_removal(image_tensor, mask) augmented = background_replacement(augmented, alt_background, aug_mask) augmented = selective_color_distort(augmented, aug_mask) return augmented
def __init__(self, args, capacity): super(ReplaySequenceMemory, self).__init__(args, capacity) if args.augment: random_shift = nn.Sequential(nn.ReplicationPad2d(4), aug.RandomCrop((84, 84))) intensity = Intensity(scale=0.1) self.augment = nn.Sequential(random_shift, intensity) else: self.augment = nn.Identity()
def test_random_crops_and_flips(self, device, dtype): width, height = 100, 100 crop_width, crop_height = 3, 3 input = torch.randn(3, 3, width, height, device=device, dtype=dtype) bbox = torch.tensor([[[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1) aug = K.AugmentationSequential( K.RandomCrop((crop_width, crop_height), padding=1, cropping_mode='resample', fill=0), K.RandomHorizontalFlip(p=1.0), data_keys=["input", "bbox_xyxy"], ) reproducibility_test((input, bbox), aug) _params = aug.forward_parameters(input.shape) # specifying the crop locations allows us to compute by hand the expected outputs crop_locations = torch.tensor( [[1.0, 2.0], [1.0, 1.0], [2.0, 0.0]], device=_params[0].data['src'].device, dtype=_params[0].data['src'].dtype, ) crops = crop_locations.expand(4, -1, -1).permute(1, 0, 2).clone() crops[:, 1:3, 0] += crop_width - 1 crops[:, 2:4, 1] += crop_height - 1 _params[0].data['src'] = crops # expected output bboxes after crop for specified crop locations and crop size (3,3) expected_out_bbox = torch.tensor( [ [[1.0, 0.0, 2.0, 1.0], [0.0, -1.0, 1.0, 1.0], [0.0, -1.0, 2.0, 0.0]], [[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]], [[0.0, 2.0, 1.0, 3.0], [-1.0, 1.0, 0.0, 3.0], [-1.0, 1.0, 1.0, 2.0]], ], device=device, dtype=dtype, ) # horizontally flip boxes based on crop width xmins = expected_out_bbox[..., 0].clone() xmaxs = expected_out_bbox[..., 2].clone() expected_out_bbox[..., 0] = crop_width - xmaxs expected_out_bbox[..., 2] = crop_width - xmins out = aug(input, bbox, params=_params) assert out[1].shape == bbox.shape assert_close(out[1], expected_out_bbox, atol=1e-4, rtol=1e-4) out_inv = aug.inverse(*out) assert out_inv[1].shape == bbox.shape assert_close(out_inv[1], bbox, atol=1e-4, rtol=1e-4)
def __init__(self, opt): super(PostTensorTransform, self).__init__() self.random_crop = ProbTransform(A.RandomCrop( (opt.input_height, opt.input_width), padding=opt.random_crop), p=0.8) self.random_rotation = ProbTransform(A.RandomRotation( opt.random_rotation), p=0.5) if opt.dataset == "cifar10": self.random_horizontal_flip = A.RandomHorizontalFlip(p=0.5)
def __init__(self, im_size=224, device=torch.device('cuda:0')): super().__init__() self.mean = torch.tensor([0.485, 0.456, 0.406]).to(device) self.std = torch.tensor([0.229, 0.224, 0.225]).to(device) self.aug = torch.nn.Sequential( kornia.geometry.transform.Resize(int(im_size * 1.2)), Kaug.RandomCrop((im_size, im_size), padding=8), Kaug.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4), Kaug.RandomHorizontalFlip(), Kaug.Normalize(mean=self.mean, std=self.std))
def test_individual_forward_and_inverse(self, device, dtype): inp = torch.randn(1, 3, 1000, 500, device=device, dtype=dtype) bbox = torch.tensor([[[355, 10], [660, 10], [660, 250], [355, 250]]], device=device, dtype=dtype) keypoints = torch.tensor([[[465, 115], [545, 116]]], device=device, dtype=dtype) mask = bbox_to_mask( torch.tensor([[[155, 0], [900, 0], [900, 400], [155, 400]]], device=device, dtype=dtype), 500, 1000)[:, None].float() crop_size = (200, 200) aug = K.AugmentationSequential( K.ImageSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.AugmentationSequential(K.ColorJiggle(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0)), K.RandomAffine(360, p=1.0), K.RandomCrop(crop_size, padding=1, cropping_mode='resample', fill=0), data_keys=['input', 'mask', 'bbox', 'keypoints'], ) reproducibility_test((inp, mask, bbox, keypoints), aug) out = aug(inp, mask, bbox, keypoints) assert out[0].shape == (*inp.shape[:2], *crop_size) assert out[1].shape == (*mask.shape[:2], *crop_size) assert out[2].shape == bbox.shape assert out[3].shape == keypoints.shape out_inv = aug.inverse(*out) assert out_inv[0].shape == inp.shape assert out_inv[1].shape == mask.shape assert out_inv[2].shape == bbox.shape assert out_inv[3].shape == keypoints.shape aug = K.AugmentationSequential(K.RandomAffine(360, p=1.0)) assert aug(inp, data_keys=['input']).shape == inp.shape aug = K.AugmentationSequential(K.RandomAffine(360, p=1.0)) assert aug(inp, data_keys=['input']).shape == inp.shape assert aug(mask, data_keys=['mask'], params=aug._params).shape == mask.shape assert aug.inverse(inp, data_keys=['input']).shape == inp.shape assert aug.inverse(bbox, data_keys=['bbox']).shape == bbox.shape assert aug.inverse(keypoints, data_keys=['keypoints']).shape == keypoints.shape assert aug.inverse(mask, data_keys=['mask']).shape == mask.shape
def __init__(self, celeba_folder, mask_folder, mode="train", use_transforms=False): super().__init__() self.images = [] self.labels = [] self.mode = mode self.mask_folder = mask_folder self.crop_aug = K.RandomCrop((218, 178), pad_if_needed=True) self.flip_aug = K.RandomHorizontalFlip() self.id_map = {} with open(celeba_folder + "/identity_CelebA.txt") as id_f: for line in id_f: im, id = line.split() self.id_map[im[:-4]] = int(id) with open(celeba_folder + "/labels/list_attr_celeba.txt") as label_f: id_count = label_f.readline().strip("\n") self.attributes = np.array(label_f.readline().strip("\n").split()) for i, line in enumerate(label_f.read().split("\n")[:-1]): image_data = line.split() image = image_data[0] if not os.path.exists(f"{mask_folder}" f"/id-{self.id_map[image[:-4]]}"): continue image_labels = [int(label) for label in image_data[1:]] image_n = int(image[:-4]) train_cond = mode == "train" and image_n < 162771 val_cond = mode == "val" and image_n >= 162771 and image_n < 182638 test_cond = mode == "test" and image_n >= 182638 full_cond = mode == "full" if train_cond or val_cond or test_cond or full_cond: self.images.append(celeba_folder + "/images/" + image) self.labels.append(image_labels) self.alt_background = TF.to_tensor(Image.open("unicorn.jpg")) self.length = len(self.images)
def __init__(self, N_TFMS: int, MAGN: int, mean: Union[tuple, list, torch.tensor], std: Union[tuple, list, torch.tensor], transform_list: list = None, use_resize: int = None, image_size: tuple = None, use_mix: int = None, mix_p: float = .5): super().__init__() self.N_TFMS, self.MAGN = N_TFMS, MAGN self.use_mix, self.mix_p = use_mix, mix_p self.image_size = image_size if not isinstance(mean, torch.Tensor): mean = torch.Tensor(mean) if not isinstance(std, torch.Tensor): std = torch.Tensor(std) if self.use_mix is not None: self.mix_list = [ K.RandomCutMix(self.image_size[0], self.image_size[1], p=1), K.RandomMixUp(p=1) ] self.use_resize = use_resize if use_resize is not None: assert len( image_size ) == 2, 'Invalid `image_size`. Must be a tuple of form (h, w)' self.resize_list = [ K.RandomResizedCrop(image_size), K.RandomCrop(image_size), K.CenterCrop(image_size) ] if self.use_resize < 3: self.resize = self.resize_list[use_resize] self.normalize = K.Normalize(mean, std) self.transform_list = transform_list if transform_list is None: self.transform_list = kornia_list(MAGN)
# # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. # ============================================================================== from __future__ import division import os import numpy as np import torch from torch import optim from torch.nn.utils import clip_grad_norm_ import kornia.augmentation as aug import torch.nn as nn from model import DQN random_shift = nn.Sequential(aug.RandomCrop((96, 46)), nn.ReplicationPad2d(4), aug.RandomCrop((100, 50))) aug = random_shift class Agent(): def __init__(self, args, env): self.args = args self.action_space = env.action_space() self.atoms = args.atoms self.Vmin = args.V_min self.Vmax = args.V_max self.support = torch.linspace(args.V_min, args.V_max, self.atoms).to(device=args.device) # Support (range) of z self.delta_z = (args.V_max - args.V_min) / (self.atoms - 1) self.batch_size = args.batch_size self.n = args.multi_step self.discount = args.discount self.norm_clip = args.norm_clip
def _setup(self, x: torch.Tensor) -> None: height, width = x.shape[-2:] self._operation = nn.Sequential( nn.ReplicationPad2d(self._shift_size), aug.RandomCrop((height, width)), )
class TestVideoSequential: @pytest.mark.parametrize('shape', [(3, 4), (2, 3, 4), (2, 3, 5, 6), (2, 3, 4, 5, 6, 7)]) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) def test_exception(self, shape, data_format, device, dtype): aug_list = K.VideoSequential(K.ColorJitter(0.1, 0.1, 0.1, 0.1), data_format=data_format, same_on_frame=True) with pytest.raises(AssertionError): img = torch.randn(*shape, device=device, dtype=dtype) aug_list(img) @pytest.mark.parametrize( 'augmentation', [ K.RandomAffine(360, p=1.0), K.CenterCrop((3, 3), p=1.0), K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomCrop((5, 5), p=1.0), K.RandomErasing(p=1.0), K.RandomGrayscale(p=1.0), K.RandomHorizontalFlip(p=1.0), K.RandomVerticalFlip(p=1.0), K.RandomPerspective(p=1.0), K.RandomResizedCrop((5, 5), p=1.0), K.RandomRotation(360.0, p=1.0), K.RandomSolarize(p=1.0), K.RandomPosterize(p=1.0), K.RandomSharpness(p=1.0), K.RandomEqualize(p=1.0), K.RandomMotionBlur(3, 35.0, 0.5, p=1.0), K.Normalize(torch.tensor([0.5, 0.5, 0.5]), torch.tensor([0.5, 0.5, 0.5]), p=1.0), K.Denormalize(torch.tensor([0.5, 0.5, 0.5]), torch.tensor([0.5, 0.5, 0.5]), p=1.0), ], ) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) def test_augmentation(self, augmentation, data_format, device, dtype): input = torch.randint(255, (1, 3, 3, 5, 6), device=device, dtype=dtype).repeat(2, 1, 1, 1, 1) / 255.0 torch.manual_seed(21) aug_list = K.VideoSequential(augmentation, data_format=data_format, same_on_frame=True) reproducibility_test(input, aug_list) @pytest.mark.parametrize( 'augmentations', [ [ K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0) ], [ K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0) ], [K.RandomAffine(360, p=1.0), kornia.color.BgrToRgb()], [ K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=0.0), K.RandomAffine(360, p=0.0) ], [K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=0.0)], [K.RandomAffine(360, p=0.0)], [ K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0), K.RandomAffine(360, p=1.0), K.RandomMixUp(p=1.0) ], ], ) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) @pytest.mark.parametrize('random_apply', [1, (1, 1), (1, ), 10, True, False]) def test_same_on_frame(self, augmentations, data_format, random_apply, device, dtype): aug_list = K.VideoSequential(*augmentations, data_format=data_format, same_on_frame=True, random_apply=random_apply) if data_format == 'BCTHW': input = torch.randn(2, 3, 1, 5, 6, device=device, dtype=dtype).repeat(1, 1, 4, 1, 1) output = aug_list(input) if aug_list.return_label: output, _ = output assert (output[:, :, 0] == output[:, :, 1]).all() assert (output[:, :, 1] == output[:, :, 2]).all() assert (output[:, :, 2] == output[:, :, 3]).all() if data_format == 'BTCHW': input = torch.randn(2, 1, 3, 5, 6, device=device, dtype=dtype).repeat(1, 4, 1, 1, 1) output = aug_list(input) if aug_list.return_label: output, _ = output assert (output[:, 0] == output[:, 1]).all() assert (output[:, 1] == output[:, 2]).all() assert (output[:, 2] == output[:, 3]).all() reproducibility_test(input, aug_list) @pytest.mark.parametrize( 'augmentations', [ [K.RandomAffine(360, p=1.0)], [K.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0)], [ K.RandomAffine(360, p=0.0), K.ImageSequential(K.RandomAffine(360, p=0.0)) ], ], ) @pytest.mark.parametrize('data_format', ["BCTHW", "BTCHW"]) def test_against_sequential(self, augmentations, data_format, device, dtype): aug_list_1 = K.VideoSequential(*augmentations, data_format=data_format, same_on_frame=False) aug_list_2 = torch.nn.Sequential(*augmentations) if data_format == 'BCTHW': input = torch.randn(2, 3, 1, 5, 6, device=device, dtype=dtype).repeat(1, 1, 4, 1, 1) if data_format == 'BTCHW': input = torch.randn(2, 1, 3, 5, 6, device=device, dtype=dtype).repeat(1, 4, 1, 1, 1) torch.manual_seed(0) output_1 = aug_list_1(input) torch.manual_seed(0) if data_format == 'BCTHW': input = input.transpose(1, 2) output_2 = aug_list_2(input.reshape(-1, 3, 5, 6)) output_2 = output_2.view(2, 4, 3, 5, 6) if data_format == 'BCTHW': output_2 = output_2.transpose(1, 2) assert (output_1 == output_2).all(), dict(aug_list_1._params) @pytest.mark.jit @pytest.mark.skip(reason="turn off due to Union Type") def test_jit(self, device, dtype): B, C, D, H, W = 2, 3, 5, 4, 4 img = torch.ones(B, C, D, H, W, device=device, dtype=dtype) op = K.VideoSequential(K.ColorJitter(0.1, 0.1, 0.1, 0.1), same_on_frame=True) op_jit = torch.jit.script(op) assert_close(op(img), op_jit(img))
def __init__(self, algo_params, env, transition_tuple=None, path=None, seed=-1): # environment self.env = PixelPybulletGym( env, image_size=algo_params['image_resize_size'], crop_size=algo_params['image_crop_size']) self.frame_stack = algo_params['frame_stack'] self.env = FrameStack(self.env, k=self.frame_stack) self.env.seed(seed) obs = self.env.reset() algo_params.update({ 'state_shape': obs.shape, # make sure the shape is like (C, H, W), not (H, W, C) 'action_dim': self.env.action_space.shape[0], 'action_max': self.env.action_space.high, 'action_scaling': self.env.action_space.high[0], }) # training args self.max_env_step = algo_params['max_env_step'] self.testing_gap = algo_params['testing_gap'] self.testing_episodes = algo_params['testing_episodes'] self.saving_gap = algo_params['saving_gap'] super(SACDrQ, self).__init__(algo_params, transition_tuple=transition_tuple, image_obs=True, training_mode='step_based', path=path, seed=seed) # torch self.encoder = PixelEncoder(self.state_shape) self.encoder_target = PixelEncoder(self.state_shape) self.network_dict.update({ 'actor': StochasticConvActor(self.action_dim, encoder=self.encoder, detach_obs_encoder=True).to(self.device), 'critic_1': ConvCritic(self.action_dim, encoder=self.encoder, detach_obs_encoder=False).to(self.device), 'critic_1_target': ConvCritic(self.action_dim, encoder=self.encoder_target, detach_obs_encoder=True).to(self.device), 'critic_2': ConvCritic(self.action_dim, encoder=self.encoder, detach_obs_encoder=False).to(self.device), 'critic_2_target': ConvCritic(self.action_dim, encoder=self.encoder_target, detach_obs_encoder=True).to(self.device), 'alpha': algo_params['alpha'], 'log_alpha': T.tensor(np.log(algo_params['alpha']), requires_grad=True, device=self.device), }) self.network_keys_to_save = ['actor'] self.actor_optimizer = Adam(self.network_dict['actor'].parameters(), lr=self.actor_learning_rate) self.critic_1_optimizer = Adam( self.network_dict['critic_1'].parameters(), lr=self.critic_learning_rate) self.critic_2_optimizer = Adam( self.network_dict['critic_2'].parameters(), lr=self.critic_learning_rate) self._soft_update(self.network_dict['critic_1'], self.network_dict['critic_1_target'], tau=1) self._soft_update(self.network_dict['critic_2'], self.network_dict['critic_2_target'], tau=1) self.target_entropy = -self.action_dim self.alpha_optimizer = Adam([self.network_dict['log_alpha']], lr=self.actor_learning_rate) # augmentation args self.image_random_shift = T.nn.Sequential( T.nn.ReplicationPad2d(4), aug.RandomCrop(self.state_shape[-2:])) self.q_regularisation_k = algo_params['q_regularisation_k'] # training args self.warmup_step = algo_params['warmup_step'] self.actor_update_interval = algo_params['actor_update_interval'] self.critic_target_update_interval = algo_params[ 'critic_target_update_interval'] # statistic dict self.statistic_dict.update({ 'episode_return': [], 'env_step_return': [], 'env_step_test_return': [], 'alpha': [], 'policy_entropy': [], })
def _setup(self, x): height, width = x.shape[-2:] self._operation = nn.Sequential(nn.ReplicationPad2d(self.shift_size), aug.RandomCrop((height, width)))
import kornia.augmentation as aug from torch import optim from torch.nn.utils import clip_grad_norm_ import torchvision.transforms as transforms import cv2 as cv # Hyper Parameters BATCH_SIZE = 32 LR = 0.01 # learning rate EPSILON = 0.9 # greedy policy GAMMA = 0.9 # reward discount TARGET_REPLACE_ITER = 100 # target update frequency MOMENTUM_UPDATE_ITER = 10 MEMORY_CAPACITY = 1000 OBSERV_MEMORY_CAPACITY = 50 random_shift = nn.Sequential(aug.RandomCrop((80, 80)), nn.ReplicationPad2d(4), aug.RandomCrop((84, 84))) img = cv.imread('/home/sam/Pictures/atari_background.jpg') transf = transforms.ToTensor() img_tensor = transf(img) R = img_tensor[0] G = img_tensor[1] B = img_tensor[2] img_tensor[0] = 0.299 * R + 0.587 * G + 0.114 * B ATARI_CANVAS = np.ceil(img_tensor[0][200:300, 500:600] * 255) class Box2DEnv(gym.Env): metadata = { 'render.modes': ['human', 'rgb_array'], 'video.frames_per_second': 2
def test_random_crops(self, device, dtype): torch.manual_seed(233) input = torch.randn(3, 3, 3, 3, device=device, dtype=dtype) bbox = torch.tensor([[[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1) points = torch.tensor([[[0.0, 0.0], [1.0, 1.0]]], device=device, dtype=dtype).expand(3, -1, -1) aug = K.AugmentationSequential( K.RandomCrop((3, 3), padding=1, cropping_mode='resample', fill=0), K.RandomAffine((360., 360.), p=1.), data_keys=["input", "mask", "bbox_xyxy", "keypoints"], ) reproducibility_test((input, input, bbox, points), aug) _params = aug.forward_parameters(input.shape) # specifying the crops allows us to compute by hand the expected outputs _params[0].data['src'] = torch.tensor( [ [[1.0, 2.0], [3.0, 2.0], [3.0, 4.0], [1.0, 4.0]], [[1.0, 1.0], [3.0, 1.0], [3.0, 3.0], [1.0, 3.0]], [[2.0, 0.0], [4.0, 0.0], [4.0, 2.0], [2.0, 2.0]], ], device=_params[0].data['src'].device, dtype=_params[0].data['src'].dtype, ) expected_out_bbox = torch.tensor( [ [[1.0, 0.0, 2.0, 1.0], [0.0, -1.0, 1.0, 1.0], [0.0, -1.0, 2.0, 0.0]], [[1.0, 1.0, 2.0, 2.0], [0.0, 0.0, 1.0, 2.0], [0.0, 0.0, 2.0, 1.0]], [[0.0, 2.0, 1.0, 3.0], [-1.0, 1.0, 0.0, 3.0], [-1.0, 1.0, 1.0, 2.0]], ], device=device, dtype=dtype, ) expected_out_points = torch.tensor( [[[0.0, -1.0], [1.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[-1.0, 1.0], [0.0, 2.0]]], device=device, dtype=dtype) out = aug(input, input, bbox, points, params=_params) assert out[0].shape == (3, 3, 3, 3) assert_close(out[0], out[1], atol=1e-4, rtol=1e-4) assert out[2].shape == bbox.shape assert_close(out[2], expected_out_bbox, atol=1e-4, rtol=1e-4) assert out[3].shape == points.shape assert_close(out[3], expected_out_points, atol=1e-4, rtol=1e-4) out_inv = aug.inverse(*out) assert out_inv[0].shape == input.shape assert_close(out_inv[0], out_inv[1], atol=1e-4, rtol=1e-4) assert out_inv[2].shape == bbox.shape assert_close(out_inv[2], bbox, atol=1e-4, rtol=1e-4) assert out_inv[3].shape == points.shape assert_close(out_inv[3], points, atol=1e-4, rtol=1e-4)
# # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. # ============================================================================== from __future__ import division import os import numpy as np import torch from torch import optim from torch.nn.utils import clip_grad_norm_ import kornia.augmentation as aug_ import torch.nn as nn from model import DQN aug_art = nn.Sequential(aug_.RandomCrop((80, 80)), nn.ReplicationPad2d(4), aug_.RandomCrop((84, 84))) aug_procgen = nn.Sequential(aug_.RandomCrop((60, 60)), nn.ReplicationPad2d(4), aug_.RandomCrop((64, 64))) class Agent(): def __init__(self, args, env): self.args = args self.action_space = env.action_space() self.atoms = args.atoms self.Vmin = args.V_min self.Vmax = args.V_max self.support = torch.linspace(args.V_min, args.V_max, self.atoms).to( device=args.device) # Support (range) of z self.delta_z = (args.V_max - args.V_min) / (self.atoms - 1)
def frost(inp, inp_coeff, frost_coeff): idx = ch.randint(len(PRECOMP_FROST), size=(len(inp), )) frost_figs = PRECOMP_FROST[idx].cuda() cropper = augmentation.RandomCrop((inp.shape[2], inp.shape[3])).cuda() return ch.clamp(inp_coeff * inp + \ frost_coeff * cropper(frost_figs), 0, 1)
def __init__(self, args=None, vis=None): super(TimeCycle, self).__init__() self.args = args if args is not None: self.kldv_coef = getattr(args, 'kldv_coef', 0) self.xent_coef = getattr(args, 'xent_coef', 0) self.zero_diagonal = getattr(args, 'zero_diagonal', 0) self.dropout_rate = getattr(args, 'dropout', 0) self.featdrop_rate = getattr(args, 'featdrop', 0) self.model_type = getattr(args, 'model_type', 'scratch') self.temperature = getattr(args, 'temp', getattr(args, 'temperature',1)) self.shuffle = getattr(args, 'shuffle', 0) self.xent_weight = getattr(args, 'xent_weight', False) else: self.kldv_coef = 0 self.xent_coef = 0 self.long_coef = 1 self.skip_coef = 0 # self.sk_align = False # self.sk_targets = True self.zero_diagonal = 0 self.dropout_rate = 0 self.featdrop_rate = 0 self.model_type = 'scratch' self.temperature = 1 self.shuffle = False self.xent_weight = False print('Model temp:', self.temperature) self.encoder = utils.make_encoder(args).cuda() self.infer_dims() self.selfsim_fc = self.make_head(depth=self.garg('head_depth', 0)) self.selfsim_head = self.make_conv3d_head(depth=1) self.context_head = self.make_conv3d_head(depth=1) # self.selfsim_head = self.make_head([self.enc_hid_dim, 2*self.enc_hid_dim, self.enc_hid_dim]) # self.context_head = self.make_head([self.enc_hid_dim, 2*self.enc_hid_dim, self.enc_hid_dim]) import resnet3d, resnet2d if self.garg('cal_coef', 0) > 0: self.stack_encoder = utils.make_stack_encoder(self.enc_hid_dim) # self.aff_encoder = resnet2d.Bottleneck(1, 128,) # # assuming no fc pre-training # for m in self.modules(): # if isinstance(m, nn.Linear): # m.weight.data.normal_(0, 0.01) # m.bias.data.zero_() self.edge = getattr(args, 'edgefunc', 'softmax') # self.kldv = torch.nn.KLDivLoss(reduction="batchmean") self.kldv = torch.nn.KLDivLoss(reduction="batchmean") self.xent = torch.nn.CrossEntropyLoss(reduction="none") self.target_temp = 1 self._xent_targets = {} self._kldv_targets = {} if self.garg('restrict', 0) > 0: self.restrict = utils.RestrictAttention(int(args.restrict)) else: self.restrict = None self.dropout = torch.nn.Dropout(p=self.dropout_rate, inplace=False) self.featdrop = torch.nn.Dropout(p=self.featdrop_rate, inplace=False) self.viz = visdom.Visdom(port=8095, env='%s_%s' % (getattr(args, 'name', 'test'), '')) #int(time.time()))) self.viz.close() if not self.viz.check_connection(): self.viz = None if vis is not None: self._viz = vis p_sz, stride = 64, 32 self.k_patch = nn.Sequential( K.RandomResizedCrop(size=(p_sz, p_sz), scale=(0.7, 0.9), ratio=(0.7, 1.3)) ) mmm, sss = torch.Tensor([0.485, 0.456, 0.406]), torch.Tensor([0.229, 0.224, 0.225]) self.k_frame = nn.Sequential( # kornia.color.Normalize(mean=-mmm/sss, std=1/sss), # K.ColorJitter(0.1, 0.1, 0.1, 0), # K.RandomResizedCrop(size=(256, 256), scale=(0.8, 0.9), ratio=(0.7, 1.3)), # kornia.color.Normalize(mean=mmm, std=sss) ) # self.k_frame_same = nn.Sequential( # K.RandomResizedCrop(size=(256, 256), scale=(0.8, 1.0), same_on_batch=True) # ) # self.k_frame_same = None self.k_frame_same = nn.Sequential( kornia.geometry.transform.Resize(256 + 20), K.RandomHorizontalFlip(same_on_batch=True), K.RandomCrop((256, 256), same_on_batch=True), ) self.unfold = torch.nn.Unfold((p_sz,p_sz), dilation=1, padding=0, stride=(stride, stride)) self.ent_stats = utils.RunningStats(1000)
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. # ============================================================================== from __future__ import division import os import numpy as np import torch from torch import optim from torch.nn.utils import clip_grad_norm_ import kornia.augmentation as aug import torch.nn as nn import torch.nn.functional as F from model.BYOLmodel import DQN random_shift = nn.Sequential(aug.RandomCrop((80, 80)), nn.ReplicationPad2d(4), aug.RandomCrop((84, 84))) aug = random_shift def D(p, z, version='simplified'): if version == 'original': z = z.detach() p = F.normalize(p, dim=1) z = F.normalize(z, dim=1) return -(p * z).sum(dim=1).mean() else: return -F.cosine_similarity(p, z.detach(), dim=-1).mean() class AgentByol():
def get_augmenter(augmenter_type: str, image_size: ImageSizeType, dataset_mean: DatasetStatType, dataset_std: DatasetStatType, padding: PaddingInputType = 1. / 8., pad_if_needed: bool = False, subset_size: int = 2) -> Union[Module, Callable]: """ Args: augmenter_type: augmenter type image_size: (height, width) image size dataset_mean: dataset mean value in CHW dataset_std: dataset standard deviation in CHW padding: percent of image size to pad on each border of the image. If a sequence of length 4 is provided, it is used to pad left, top, right, bottom borders respectively. If a sequence of length 2 is provided, it is used to pad left/right, top/bottom borders, respectively. pad_if_needed: bool flag for RandomCrop "pad_if_needed" option subset_size: number of augmentations used in subset Returns: nn.Module for Kornia augmentation or Callable for torchvision transform """ if not isinstance(padding, tuple): assert isinstance(padding, float) padding = (padding, padding, padding, padding) assert len(padding) == 2 or len(padding) == 4 if len(padding) == 2: # padding of length 2 is used to pad left/right, top/bottom borders, respectively # padding of length 4 is used to pad left, top, right, bottom borders respectively padding = (padding[0], padding[1], padding[0], padding[1]) # image_size is of shape (h,w); padding values is [left, top, right, bottom] borders padding = (int(image_size[1] * padding[0]), int( image_size[0] * padding[1]), int(image_size[1] * padding[2]), int(image_size[0] * padding[3])) augmenter_type = augmenter_type.strip().lower() if augmenter_type == "simple": return nn.Sequential( K.RandomCrop(size=image_size, padding=padding, pad_if_needed=pad_if_needed, padding_mode='reflect'), K.RandomHorizontalFlip(p=0.5), K.Normalize(mean=torch.tensor(dataset_mean, dtype=torch.float32), std=torch.tensor(dataset_std, dtype=torch.float32)), ) elif augmenter_type == "fixed": return nn.Sequential( K.RandomHorizontalFlip(p=0.5), # K.RandomVerticalFlip(p=0.2), K.RandomResizedCrop(size=image_size, scale=(0.8, 1.0), ratio=(1., 1.)), RandomAugmentation(p=0.5, augmentation=F.GaussianBlur2d( kernel_size=(3, 3), sigma=(1.5, 1.5), border_type='constant')), K.ColorJitter(contrast=(0.75, 1.5)), # additive Gaussian noise K.RandomErasing(p=0.1), # Multiply K.RandomAffine(degrees=(-25., 25.), translate=(0.2, 0.2), scale=(0.8, 1.2), shear=(-8., 8.)), K.Normalize(mean=torch.tensor(dataset_mean, dtype=torch.float32), std=torch.tensor(dataset_std, dtype=torch.float32)), ) elif augmenter_type in ["validation", "test"]: return nn.Sequential( K.Normalize(mean=torch.tensor(dataset_mean, dtype=torch.float32), std=torch.tensor(dataset_std, dtype=torch.float32)), ) elif augmenter_type == "randaugment": return nn.Sequential( K.RandomCrop(size=image_size, padding=padding, pad_if_needed=pad_if_needed, padding_mode='reflect'), K.RandomHorizontalFlip(p=0.5), RandAugmentNS(n=subset_size, m=10), K.Normalize(mean=torch.tensor(dataset_mean, dtype=torch.float32), std=torch.tensor(dataset_std, dtype=torch.float32)), ) else: raise NotImplementedError( f"\"{augmenter_type}\" is not a supported augmenter type")