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")
