def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)):
super(Residual, self).__init__()
modules = []
for _ in range(num_block):
modules.append(Depth_Wise(c, c, residual=True, kernel=kernel, padding=padding, stride=stride, groups=groups))
self.model = Sequential(*modules)
def check(broadcast_buffers: bool,
grad_accumulation: bool = False) -> None:
# Any model works. Add one different buffer per rank
model = Sequential(Linear(2, 3), Linear(3, 3), Linear(3, 3),
Linear(3, 3), Linear(3, 3), Linear(3, 3))
model.register_buffer("test_buffer", torch.ones((1)) * rank)
model.to(device)
next(model.parameters()
).requires_grad = False # Test non-trainable parameters
optimizer = OSS(params=model.parameters(),
optim=torch.optim.SGD,
lr=0.01,
momentum=0.99)
ddp_model = ShardedDataParallel(model,
optimizer,
broadcast_buffers=broadcast_buffers)
def check_same_model_params(same_params: bool):
# Check that all the params are the same on all ranks
# This should be true with and without broadcast_buffers, we don't have any real buffer here
receptacle: List[torch.Tensor] = []
if dist.get_backend() != "nccl":
for pg in optimizer.param_groups:
for p in pg["params"]:
# Check the params
receptacle = [p.clone() for _ in range(world_size)
] if rank == 0 else []
dist.gather(p, receptacle, dst=0)
if rank == 0:
for sync_p in receptacle[1:]:
if same_params:
assert torch.all(
torch.eq(receptacle[0], sync_p)
), "Models differ in between ranks"
else:
assert not torch.all(
torch.eq(receptacle[0], sync_p)
), "Gradients should not have been synced"
# Check that all the buffers are in sync (authoritative rank is 0, its buffer is 0)
if broadcast_buffers:
for b in ddp_model.buffers():
receptacle = [b.clone() for _ in range(world_size)
] if rank == 0 else []
dist.gather(b, receptacle, dst=0)
if rank == 0:
for sync_b in receptacle[1:]:
if same_params:
assert torch.all(
torch.eq(receptacle[0], sync_b)
), "Models differ in between ranks"
else:
assert not torch.all(
torch.eq(receptacle[0], sync_b)
), "Gradients should not have been synced"
assert b.cpu().item() == 0.0
# The model should be synchronized in between the ranks at ShardedDataParallel construction time, check that
check_same_model_params(same_params=True)
# Optim loop
def closure():
optimizer.zero_grad()
with ddp_model.no_sync() if grad_accumulation else suppress():
input_tensor = torch.rand((64, 2)).to(device)
loss = ddp_model(input_tensor).abs().sum()
loss.backward()
return loss
# The models should stay the same in between the ranks
for i in range(5):
_ = optimizer.step(closure=closure)
# when running on cpu/gloo the "nodes" are not really different
same_params = device == torch.device("cpu") or grad_accumulation
check_same_model_params(same_params=same_params)
def create_vgg_ssd(num_classes, is_test=False):
vgg_config = [
64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
512, 512, 512
]
base_net = ModuleList(vgg(vgg_config))
source_layer_indexes = [
(23, BatchNorm2d(512)),
len(base_net),
]
extras = ModuleList([
Sequential(
Conv2d(in_channels=1024, out_channels=256, kernel_size=1), ReLU(),
Conv2d(in_channels=256,
out_channels=512,
kernel_size=3,
stride=2,
padding=1), ReLU()),
Sequential(
Conv2d(in_channels=512, out_channels=128, kernel_size=1), ReLU(),
Conv2d(in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1), ReLU()),
Sequential(Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(in_channels=128, out_channels=256, kernel_size=3),
ReLU()),
Sequential(Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(in_channels=128, out_channels=256, kernel_size=3),
ReLU())
])
regression_headers = ModuleList([
Conv2d(in_channels=512, out_channels=4 * 4, kernel_size=3, padding=1),
Conv2d(in_channels=1024, out_channels=6 * 4, kernel_size=3, padding=1),
Conv2d(in_channels=512, out_channels=6 * 4, kernel_size=3, padding=1),
Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=3, padding=1),
Conv2d(in_channels=256, out_channels=4 * 4, kernel_size=3, padding=1),
Conv2d(in_channels=256, out_channels=4 * 4, kernel_size=3,
padding=1), # TODO: change to kernel_size=1, padding=0?
])
classification_headers = ModuleList([
Conv2d(in_channels=512,
out_channels=4 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=1024,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=512,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=256,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=256,
out_channels=4 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=256,
out_channels=4 * num_classes,
kernel_size=3,
padding=1), # TODO: change to kernel_size=1, padding=0?
])
return SSD(num_classes,
base_net,
source_layer_indexes,
extras,
classification_headers,
regression_headers,
is_test=is_test,
config=config)
def __init__(self):
super().__init__()
self.layer = Sequential(
OrderedDict([("mlp_1", nn.Linear(32, 32)),
("mlp_2", nn.Linear(32, 32, bias=False)),
("mlp_3", nn.Linear(32, 2))]))
def __init__(
self,
nb_input_channels,
board_height,
board_width,
channels=256,
channels_operating_init=224,
channel_expansion=32,
act_type='relu',
channels_value_head=8,
channels_policy_head=81,
value_fc_size=256,
dropout_rate=0.15,
select_policy_from_plane=True,
kernels=None,
n_labels=4992,
se_types=None,
use_avg_features=False,
use_wdl=False,
use_plys_to_end=False,
use_mlp_wdl_ply=False,
bn_mom=0.9,
):
"""
RISEv3 architecture
:param channels: Main number of channels
:param channels_operating: Initial number of channels at the start of the net for the depthwise convolution
:param channel_expansion: Number of channels to add after each residual block
:param act_type: Activation type to use
:param channels_value_head: Number of channels for the value head
:param value_fc_size: Number of units in the fully connected layer of the value head
:param channels_policy_head: Number of channels for the policy head
:param dropout_rate: Droput factor to use. If 0, no dropout will be applied. Value must be in [0,1]
:param select_policy_from_plane: True, if policy head type shall be used
:param kernels: List of kernel sizes used for the residual blocks. The length of the list corresponds to the number
of residual blocks.
:param n_labels: Number of policy target labels (used for select_policy_from_plane=False)
:param se_types: List of squeeze exciation modules to use for each residual layer.
The length of this list must be the same as len(kernels). Available types:
- "se": Squeeze excitation block - Hu et al. - https://arxiv.org/abs/1709.01507
- "cbam": Convolutional Block Attention Module (CBAM) - Woo et al. - https://arxiv.org/pdf/1807.06521.pdf
- "ca_se": Same as "se"
- "cm_se": Squeeze excitation with max operator
- "sa_se": Spatial excitation with average operator
- "sm_se": Spatial excitation with max operator
the spatial dimensionality and the number of channels will be doubled.
Later the spatial and scalar embeddings will be merged again.
:param use_wdl: If a win draw loss head shall be used
:param use_plys_to_end: If a plys to end prediction head shall be used
:param use_mlp_wdl_ply: If a small mlp with value output for the wdl and ply head shall be used
:return: symbol
"""
super(RiseV3, self).__init__()
self.nb_input_channels = nb_input_channels
self.use_plys_to_end = use_plys_to_end
self.use_wdl = use_wdl
if len(kernels) != len(se_types):
raise Exception(
f'The length of "kernels": {len(kernels)} must be the same as'
f' the length of "se_types": {len(se_types)}')
valid_se_types = [
None, "se", "cbam", "eca_se", "ca_se", "cm_se", "sa_se", "sm_se"
]
for se_type in se_types:
if se_type not in valid_se_types:
raise Exception(
f"Unavailable se_type: {se_type}. Available se_types include {se_types}"
)
res_blocks = _get_res_blocks(act_type, bn_mom, channels,
channels_operating_init,
channel_expansion, kernels, se_types)
self.body_spatial = Sequential(
_Stem(channels=channels,
bn_mom=bn_mom,
act_type=act_type,
nb_input_channels=nb_input_channels),
*res_blocks,
)
self.nb_body_spatial_out = channels * board_height * board_width
# create the two heads which will be used in the hybrid fwd pass
self.value_head = _ValueHead(board_height, board_width, channels,
channels_value_head, value_fc_size,
bn_mom, act_type, False,
nb_input_channels, use_wdl,
use_plys_to_end, use_mlp_wdl_ply)
self.policy_head = _PolicyHead(board_height, board_width, channels,
channels_policy_head, n_labels, bn_mom,
act_type, select_policy_from_plane)
def create_squeezenet_ssd_lite(num_classes, is_test=False):
base_net = squeezenet1_1(False).features # disable dropout layer
source_layer_indexes = [12]
extras = ModuleList([
Sequential(
Conv2d(in_channels=512, out_channels=256, kernel_size=1),
ReLU(),
SeperableConv2d(in_channels=256,
out_channels=512,
kernel_size=3,
stride=2,
padding=2),
),
Sequential(
Conv2d(in_channels=512, out_channels=256, kernel_size=1),
ReLU(),
SeperableConv2d(in_channels=256,
out_channels=512,
kernel_size=3,
stride=2,
padding=1),
),
Sequential(
Conv2d(in_channels=512, out_channels=128, kernel_size=1),
ReLU(),
SeperableConv2d(in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1),
),
Sequential(
Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
SeperableConv2d(in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1),
),
Sequential(
Conv2d(in_channels=256, out_channels=128, kernel_size=1), ReLU(),
SeperableConv2d(in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1))
])
regression_headers = ModuleList([
SeperableConv2d(in_channels=512,
out_channels=6 * 4,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=512,
out_channels=6 * 4,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=512,
out_channels=6 * 4,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=256,
out_channels=6 * 4,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=256,
out_channels=6 * 4,
kernel_size=3,
padding=1),
Conv2d(in_channels=256, out_channels=6 * 4, kernel_size=1),
])
classification_headers = ModuleList([
SeperableConv2d(in_channels=512,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=512,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=512,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=256,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=256,
out_channels=6 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=256, out_channels=6 * num_classes, kernel_size=1),
])
return SSD(num_classes,
base_net,
source_layer_indexes,
extras,
classification_headers,
regression_headers,
is_test=is_test,
config=config)
def __init__(self, drop_ratio=0.6):
super(PreheadID, self).__init__()
self.fc = Sequential(BatchNorm2d(512), Dropout(drop_ratio), Flatten(),
Linear(512 * 7 * 7, 512), BatchNorm1d(512))
def __init__(self) -> None:
super().__init__()
self.trunk = Sequential(Linear(2, 3), Linear(3, 3), Linear(3, 3))
self.head = Sequential(Linear(3, 3), Linear(3, 3))
self.multiple_fw = multiple_fw
self.embedding = torch.nn.Embedding(EMB_SIZE, 2)
def __init__(self, classnum):
super(Backbone, self).__init__()
unit_module = bottleneck_IR
self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
BatchNorm2d(64,momentum=0.9,eps=2e-5),
PReLU(64))
self.output_layer = Sequential(BatchNorm2d(512,momentum=0.9,eps=2e-5),
Dropout(0.4),
Flatten(),
Linear(512 * 7 * 7, 512),
BatchNorm1d(512,momentum=0.9,eps=2e-5))
self.layer1=make_layers(unit_module,64,64,3)
self.layer2_1=make_layers(unit_module,64,128,6)
self.layer2_2=make_layers(unit_module,128,128,7,stride=1)
self.layer3_1=make_layers(unit_module,128,256,10)
self.layer3_2=make_layers(unit_module,256,256,10,stride=1)
self.layer3_3=make_layers(unit_module,256,256,10,stride=1)
self.layer4=make_layers(unit_module,256,512,3)
self.ac_fc=Arcface(embedding_size=512,classnum=classnum)
self.att1=Sequential(
Conv2d(3, 64, (3, 3), 2, 1, bias=False),
BatchNorm2d(64,momentum=0.9,eps=2e-5),
PReLU(64),
Conv2d(64, 64, (3, 3), 2, 1, bias=False),
BatchNorm2d(64, momentum=0.9, eps=2e-5),
PReLU(64),
Conv2d(64, 64, (3, 3), 2, 1, bias=False),
BatchNorm2d(64, momentum=0.9, eps=2e-5),
PReLU(64),
Conv2d(64, 64, (3, 3), 2, 1, bias=False),
BatchNorm2d(64, momentum=0.9, eps=2e-5),
PReLU(64),
AdaptiveAvgPool2d(1),
Conv2d(64, 64, kernel_size=1, padding=0, bias=False),
BatchNorm2d(64, momentum=0.9, eps=2e-5),
Sigmoid()
)
self.att2_1=Sequential(
Conv2d(64, 128, (3, 3), 2, 1, bias=False),
BatchNorm2d(128,momentum=0.9,eps=2e-5),
PReLU(128),
Conv2d(128, 128, (3, 3), 2, 1, bias=False),
BatchNorm2d(128, momentum=0.9, eps=2e-5),
PReLU(128),
Conv2d(128, 128, (3, 3), 2, 1, bias=False),
BatchNorm2d(128, momentum=0.9, eps=2e-5),
PReLU(128),
AdaptiveAvgPool2d(1),
Conv2d(128, 128, kernel_size=1, padding=0, bias=False),
BatchNorm2d(128, momentum=0.9, eps=2e-5),
Sigmoid()
)
self.att2_2=Sequential(
Conv2d(128, 128, (3, 3), 2, 1, bias=False),
BatchNorm2d(128,momentum=0.9,eps=2e-5),
PReLU(128),
Conv2d(128, 128, (3, 3), 2, 1, bias=False),
BatchNorm2d(128, momentum=0.9, eps=2e-5),
PReLU(128),
Conv2d(128, 128, (3, 3), 1, 1, bias=False),
BatchNorm2d(128, momentum=0.9, eps=2e-5),
PReLU(128),
AdaptiveAvgPool2d(1),
Conv2d(128, 128, kernel_size=1, padding=0, bias=False),
BatchNorm2d(128, momentum=0.9, eps=2e-5),
Sigmoid()
)
self.att3_1=Sequential(
Conv2d(128, 256, (3, 3), 2, 1, bias=False),
BatchNorm2d(256,momentum=0.9,eps=2e-5),
PReLU(256),
Conv2d(256, 256, (3, 3), 2, 1, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
PReLU(256),
Conv2d(256, 256, (3, 3), 1, 1, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
PReLU(256),
AdaptiveAvgPool2d(1),
Conv2d(256, 256, kernel_size=1, padding=0, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
Sigmoid()
)
self.att3_2=Sequential(
Conv2d(256, 256, (3, 3), 2, 1, bias=False),
BatchNorm2d(256,momentum=0.9,eps=2e-5),
PReLU(256),
Conv2d(256, 256, (3, 3), 1, 1, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
PReLU(256),
Conv2d(256, 256, (3, 3), 1, 1, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
PReLU(256),
AdaptiveAvgPool2d(1),
Conv2d(256, 256, kernel_size=1, padding=0, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
Sigmoid()
)
self.att3_3=Sequential(
Conv2d(256, 256, (3, 3), 2, 1, bias=False),
BatchNorm2d(256,momentum=0.9,eps=2e-5),
PReLU(256),
Conv2d(256, 256, (3, 3), 1, 1, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
PReLU(256),
Conv2d(256, 256, (3, 3), 1, 1, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
PReLU(256),
AdaptiveAvgPool2d(1),
Conv2d(256, 256, kernel_size=1, padding=0, bias=False),
BatchNorm2d(256, momentum=0.9, eps=2e-5),
Sigmoid()
)
self.att4=Sequential(
Conv2d(256, 512, (3, 3), 2, 1, bias=False),
BatchNorm2d(512,momentum=0.9,eps=2e-5),
PReLU(512),
Conv2d(512, 512, (3, 3), 1, 1, bias=False),
BatchNorm2d(512, momentum=0.9, eps=2e-5),
PReLU(512),
Conv2d(512, 512, (3, 3), 1, 1, bias=False),
BatchNorm2d(512, momentum=0.9, eps=2e-5),
PReLU(512),
AdaptiveAvgPool2d(1),
Conv2d(512, 512, kernel_size=1, padding=0, bias=False),
BatchNorm2d(512, momentum=0.9, eps=2e-5),
Sigmoid()
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(3. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
n = m.out_features
m.weight.data.normal_(0, math.sqrt(3. / n))
def create_mb_tiny_fd(num_classes,
is_test=False,
device="cuda",
without_postprocessing=False):
base_net = Mb_Tiny(2)
base_net_model = base_net.model # disable dropout layer
source_layer_indexes = [8, 11, 13]
extras = ModuleList([
Sequential(
Conv2d(in_channels=base_net.base_channel * 16,
out_channels=base_net.base_channel * 4,
kernel_size=1), ReLU(),
SeperableConv2d(in_channels=base_net.base_channel * 4,
out_channels=base_net.base_channel * 16,
kernel_size=3,
stride=2,
padding=1), ReLU())
])
regression_headers = ModuleList([
SeperableConv2d(in_channels=base_net.base_channel * 4,
out_channels=3 * 4,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=base_net.base_channel * 8,
out_channels=2 * 4,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=base_net.base_channel * 16,
out_channels=2 * 4,
kernel_size=3,
padding=1),
Conv2d(in_channels=base_net.base_channel * 16,
out_channels=3 * 4,
kernel_size=3,
padding=1)
])
classification_headers = ModuleList([
SeperableConv2d(in_channels=base_net.base_channel * 4,
out_channels=3 * num_classes,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=base_net.base_channel * 8,
out_channels=2 * num_classes,
kernel_size=3,
padding=1),
SeperableConv2d(in_channels=base_net.base_channel * 16,
out_channels=2 * num_classes,
kernel_size=3,
padding=1),
Conv2d(in_channels=base_net.base_channel * 16,
out_channels=3 * num_classes,
kernel_size=3,
padding=1)
])
return SSD(num_classes,
base_net_model,
source_layer_indexes,
extras,
classification_headers,
regression_headers,
is_test=is_test,
config=config,
device=device,
without_postprocessing=without_postprocessing)
def test_gradcam_wiring(mocker, layers, layer_out, id_img):
"""Check gradam calculation."""
out = (
torch.ones((1, 3)),
torch.ones((1, 3, 1, 1)),
torch.ones((1, 1, 4, 4)),
torch.ones((1, 1, 4, 4)),
)
out[0].detach = mocker.Mock(return_value=out[0])
out[0].squeeze = mocker.Mock(return_value=out[0])
# Wire mocks together
mocker.patch(
"midnite.visualization.base.methods.Backpropagation.visualize",
return_value=out[0],
)
mocker.patch("torch.nn.functional.relu", return_value=out[3])
split = (mocker.Mock(spec=LayerSplit), mocker.Mock(spec=LayerSplit))
split[0].get_mean = mocker.Mock(return_value=out[2])
split[0].invert = mocker.Mock(return_value=split[1])
split[1].fill_dimensions = mocker.Mock(return_value=out[1])
mul_out = torch.zeros((1, 1, 3, 4))
layer_out[1].mul_ = mocker.Mock(return_value=mul_out)
sut = GradAM(
mocker.Mock(spec=torch.nn.Module),
mocker.Mock(spec=NeuronSelector),
Sequential(*layers),
split[0],
)
result = sut.visualize(id_img)
# Check preprocessing
id_img.clone.assert_called_once()
id_img.detach.assert_called_once()
id_img.to.assert_called_once()
layers[0].train.assert_not_called()
layers[1].train.assert_not_called()
# Check Calculation
# 1. forward pass
layers[0].assert_called_once()
layers[1].assert_called_once()
# 2. backpropagation
methods.Backpropagation.visualize.assert_called_with(layer_out[1])
# 3. Split invert and fill dimensions
split[0].invert.assert_called()
split[1].fill_dimensions(out) # -> out[1]
# 4. Multiply together
layer_out[1].mul_.assert_called_once_with(out[1]) # -> layer_out[2]
# 5. Calculate mean
split[0].get_mean.assert_called_with(mul_out)
# 6. Relu
torch.nn.functional.relu.assert_called_once_with(out[2])
# Check result
assert_that(result).is_same_as(out[3])
def __init__(self):
super().__init__()
self.layers = Sequential(FSDP(Linear(5, 5)))
def __init__(self, backbone_net):
super(AuxillaryConvolutions, self).__init__()
if backbone_net == "MobileNetV2":
self.extras = ModuleList(
[
Sequential(
Conv2d(in_channels=1280, out_channels=256, kernel_size=1),
ReLU(),
Conv2d(
in_channels=256,
out_channels=512,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
Sequential(
Conv2d(in_channels=512, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
Sequential(
Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
Sequential(
Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
]
)
self.init_conv2d()
elif backbone_net == "MobileNetV1":
self.extras = ModuleList(
[
Sequential(
Conv2d(in_channels=1024, out_channels=256, kernel_size=1),
ReLU(),
Conv2d(
in_channels=256,
out_channels=512,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
Sequential(
Conv2d(in_channels=512, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
Sequential(
Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
Sequential(
Conv2d(in_channels=256, out_channels=128, kernel_size=1),
ReLU(),
Conv2d(
in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
),
ReLU(),
),
]
)
self.init_conv2d()
def __init__(self):
super(Net, self).__init__()
#name
self.name = "DirCNNh3"
#optimizer
self.lr = 0.001
self.optimizer_name = 'Adam-Exp'
#data
#self.data_name = "ModelNet10"
self.data_name = "Geometry"
self.batch_size = 20
self.nr_points = 1024
self.nr_classes = 10 if self.data_name == 'ModelNet10' else 40
#train_info
self.max_epochs = 301
self.save_every = 100
#model
self.k = 20
self.l = 7
# DD1
self.in_size = 3
self.out_size = 64
layers = []
layers.append(Linear(self.in_size, 64))
layers.append(ReLU())
layers.append(torch.nn.BatchNorm1d(64))
layers.append(Linear(64 , 64))
layers.append(ReLU())
layers.append(torch.nn.BatchNorm1d(64))
layers.append(Linear(64, self.out_size))
layers.append(ReLU())
layers.append(torch.nn.BatchNorm1d(self.out_size))
dense3dnet = Sequential(*layers)
self.dd = DD(l = self.l,
k = self.k,
mlp = dense3dnet,
conv_p = True,
conv_fc = False,
conv_fn = False,
out_3d = True)
# DD2
self.in_size_2 = 64 * 3 + 3
self.out_size_2 = 128
layers2 = []
layers2.append(Linear(self.in_size_2, self.out_size_2))
layers2.append(ReLU())
layers2.append(torch.nn.BatchNorm1d(self.out_size_2))
dense3dnet2 = Sequential(*layers2)
self.dd2 = DD(l = self.l,
k = self.k,
mlp = dense3dnet2,
conv_p = True,
conv_fc = True,
conv_fn = True,
out_3d = False)
self.nn1 = torch.nn.Linear(self.out_size_2, 1024)
self.bn1 = torch.nn.BatchNorm1d(1024)
self.nn2 = torch.nn.Linear(1024, 512)
self.bn2 = torch.nn.BatchNorm1d(512)
self.nn3 = torch.nn.Linear(512, 265)
self.bn3 = torch.nn.BatchNorm1d(265)
self.nn4 = torch.nn.Linear(265, self.nr_classes)
self.sm = torch.nn.LogSoftmax(dim=1)
def __init__(self):
super().__init__()
self.layers = Sequential(FSDP(Linear(5, 5), **fsdp_config), )
xs = np.random.randint(low=0, high=2, size=(batch_size, 2))
ys = np.array([operator.xor(*x) for x in xs])
xs = Variable(torch.from_numpy(xs).float())
ys = Variable(torch.from_numpy(ys).long())
return xs, ys
def argmax(tensor, dim=1):
return tensor.max(dim=dim)[1]
H = 8
model = Sequential(Linear(2, H), ReLU(), Linear(H, 2))
if __name__ == "__main__":
criterion = F.cross_entropy
optimizer = optim.Adam(model.parameters())
# One batch is enough since XOR only has 4 possible values.
xs, ys = gen_data()
for i in range(N):
preds = model(xs)
loss = criterion(preds, ys)
print(float(loss))
optimizer.zero_grad()
loss.backward()
def __init__(self, h_nf):
self.node_emb_out = Sequential(Linear(h_nf, h_nf // 2), ShiftedSoftplus(), Linear(h_nf // 2, 1))
def __init__(self, n_output=1,num_features_xd=78, num_features_xt=25,
n_filters=32, embed_dim=128, output_dim=128, dropout=0.2):
super(GINConvNet, self).__init__()
dim = 32
self.dropout = nn.Dropout(dropout)
self.relu = nn.ReLU()
self.n_output = n_output
# convolution layers
nn1 = Sequential(Linear(num_features_xd, dim), ReLU(), Linear(dim, dim))
self.conv1 = GINConv(nn1)
self.bn1 = torch.nn.BatchNorm1d(dim)
nn2 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))
self.conv2 = GINConv(nn2)
self.bn2 = torch.nn.BatchNorm1d(dim)
nn3 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))
self.conv3 = GINConv(nn3)
self.bn3 = torch.nn.BatchNorm1d(dim)
nn4 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))
self.conv4 = GINConv(nn4)
self.bn4 = torch.nn.BatchNorm1d(dim)
nn5 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))
self.conv5 = GINConv(nn5)
self.bn5 = torch.nn.BatchNorm1d(dim)
self.fc1_xd = Linear(dim, output_dim)
# 1D convolution on protein sequence # mut features
self.embedding_xt_mut = nn.Embedding(num_features_xt + 1, embed_dim)
self.conv_xt_mut_1 = nn.Conv1d(in_channels=1000, out_channels=n_filters, kernel_size=8)
# 1D convolution on protein sequence # mut features
self.embedding_xt_meth = nn.Embedding(num_features_xt + 1, embed_dim)
self.conv_xt_meth_1 = nn.Conv1d(in_channels=1000, out_channels=n_filters, kernel_size=8)
# cell line mut feature
self.conv_xt_mut_1 = nn.Conv1d(in_channels=1, out_channels=n_filters, kernel_size=8)
self.pool_xt_mut_1 = nn.MaxPool1d(3)
self.conv_xt_mut_2 = nn.Conv1d(in_channels=n_filters, out_channels=n_filters*2, kernel_size=8)
self.pool_xt_mut_2 = nn.MaxPool1d(3)
self.conv_xt_mut_3 = nn.Conv1d(in_channels=n_filters*2, out_channels=n_filters*4, kernel_size=8)
self.pool_xt_mut_3 = nn.MaxPool1d(3)
self.fc1_xt_mut = nn.Linear(2944, output_dim)
# cell line meth feature
self.conv_xt_meth_1 = nn.Conv1d(in_channels=1, out_channels=n_filters, kernel_size=8)
self.pool_xt_meth_1 = nn.MaxPool1d(3)
self.conv_xt_meth_2 = nn.Conv1d(in_channels=n_filters, out_channels=n_filters*2, kernel_size=8)
self.pool_xt_meth_2 = nn.MaxPool1d(3)
self.conv_xt_meth_3 = nn.Conv1d(in_channels=n_filters*2, out_channels=n_filters*4, kernel_size=8)
self.pool_xt_meth_3 = nn.MaxPool1d(3)
self.fc1_xt_meth = nn.Linear(83584, output_dim)
# combined layers
self.fc1 = nn.Linear(3*output_dim, 1024)
self.fc2 = nn.Linear(1024, 128)
self.out = nn.Linear(128, n_output)
# activation and regularization
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.5)
def build(self, input_size: int) -> Module:
if self.variable_metadata.is_binary():
return Sequential(Linear(input_size, self.size), Sigmoid())
elif self.variable_metadata.is_numerical():
return Linear(input_size, self.size)
BatchDiagHessian,
BatchGrad,
BatchL2Grad,
DiagGGNExact,
DiagGGNMC,
DiagHessian,
SqrtGGNExact,
SqrtGGNMC,
SumGradSquared,
Variance,
)
from backpack.utils.examples import load_one_batch_mnist
X, y = load_one_batch_mnist(batch_size=512)
model = Sequential(Flatten(), Linear(784, 10))
lossfunc = CrossEntropyLoss()
model = extend(model)
lossfunc = extend(lossfunc)
# %%
# First order extensions
# ----------------------
# %%
# Batch gradients
loss = lossfunc(model(X), y)
with backpack(BatchGrad()):
loss.backward()
def __init__(self):
super(Attrhead, self).__init__()
self.fc = Sequential(BatchNorm2d(512), Flatten(),
Linear(512 * 7 * 7, 256), BatchNorm1d(256),
Linear(256, 40))
from torch import Tensor
from torch.nn.functional import mse_loss
from torch.nn import Sequential
# Imports from aihwkit.
from aihwkit.nn import AnalogLinear
from aihwkit.optim.analog_sgd import AnalogSGD
from aihwkit.simulator.devices import ConstantStepResistiveDevice
# Prepare the datasets (input and expected output).
x_b = Tensor([[0.1, 0.2, 0.0, 0.0], [0.2, 0.4, 0.0, 0.0]])
y_b = Tensor([[0.3], [0.6]])
# Define a multiple-layer network, using a constant step device type.
model = Sequential(
AnalogLinear(4, 2, resistive_device=ConstantStepResistiveDevice()),
AnalogLinear(2, 2, resistive_device=ConstantStepResistiveDevice()),
AnalogLinear(2, 1, resistive_device=ConstantStepResistiveDevice()))
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.5)
opt.regroup_param_groups(model)
for epoch in range(100):
# Add the training Tensor to the model (input).
pred = model(x_b)
# Add the expected output Tensor.
loss = mse_loss(pred, y_b)
# Run training (backward propagation).
loss.backward()
opt.step()
def __init__(self,
edge_dim,
x_dim,
u_dim,
dim=32,
bottleneck_ratio=0.25,
normalization_cls=None,
activation_cls=nn.ReLU,
dropout_cls=nn.Dropout,
dropout_prob=0.,
residual=True,
pooling='mean'):
super(MegNetBlock_v4, self).__init__()
self.dim = dim
self.residual = residual
self.pooling = pooling
kwargs = dict(normalization_cls=normalization_cls,
activation_cls=activation_cls,
dropout_cls=dropout_cls,
dropout_prob=dropout_prob)
bottleneck_dim = int(dim * bottleneck_ratio)
def one_layer_block(in_, out_):
return Sequential(*linear_block(
in_,
out_,
normalization=normalization_cls(in_)
if normalization_cls else None,
activation=activation_cls() if activation_cls else None,
dropout=dropout_cls(dropout_prob) if dropout_cls else None))
self.edge_dense = one_layer_block(edge_dim, bottleneck_dim)
self.node_dense = one_layer_block(x_dim, bottleneck_dim)
self.global_dense = one_layer_block(u_dim, bottleneck_dim)
self.edge_msg = create_mlp_v2(
input_dim=bottleneck_dim * 4,
output_dim=bottleneck_dim,
hidden_dims=[bottleneck_dim * 2, bottleneck_dim * 2],
**kwargs)
self.node_msg = create_mlp_v2(
input_dim=bottleneck_dim * 3,
output_dim=bottleneck_dim,
hidden_dims=[bottleneck_dim * 2, bottleneck_dim * 2],
**kwargs)
self.global_msg = create_mlp_v2(
input_dim=bottleneck_dim * 3,
output_dim=bottleneck_dim,
hidden_dims=[bottleneck_dim * 2, bottleneck_dim * 2],
**kwargs)
self.edge_out_dense = Sequential(
*linear_block(bottleneck_dim,
dim,
normalization=normalization_cls(bottleneck_dim)
if normalization_cls else None))
self.node_out_dense = Sequential(
*linear_block(bottleneck_dim,
dim,
normalization=normalization_cls(bottleneck_dim)
if normalization_cls else None))
self.global_out_dense = Sequential(
*linear_block(bottleneck_dim,
dim,
normalization=normalization_cls(bottleneck_dim)
if normalization_cls else None))
def __init__(self, features, classes):
super(Decoder, self).__init__()
self.layers = Sequential(
Conv1d(features, classes, kernel_size=1, bias=True))
def run_ddp_parity_two_optim(rank, world_size, backend, temp_file_name):
url = "file://" + temp_file_name
dist.init_process_group(init_method=url,
backend=backend,
rank=rank,
world_size=world_size)
device = torch.device("cuda")
torch.cuda.set_device(rank)
torch.manual_seed(rank)
np.random.seed(rank) # Any model works. Add one different buffer per rank
model = Sequential(Linear(2, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3),
Linear(3, 3), Linear(3, 3))
model.register_buffer("test_buffer", torch.ones((1)) * rank)
model.to(device)
n_half_params = len(list(model.parameters())) // 2
sharded_optimizer = OSS(params=list(model.parameters())[:n_half_params],
optim=torch.optim.SGD,
lr=1e-3,
momentum=0.99)
sharded_optimizer_2 = OSS(params=list(model.parameters())[n_half_params:],
optim=torch.optim.SGD,
lr=1e-3,
momentum=0.99)
sharded_ddp_model = ShardedDataParallel(
module=model,
sharded_optimizer=sharded_optimizer,
broadcast_buffers=True)
ddp_model_single = copy.deepcopy(model)
ddp_optimizer = torch.optim.SGD(list(
ddp_model_single.parameters())[:n_half_params],
lr=1e-3,
momentum=0.99)
ddp_optimizer_2 = torch.optim.SGD(list(
ddp_model_single.parameters())[n_half_params:],
lr=1e-3,
momentum=0.99)
ddp_model = DDP(ddp_model_single,
device_ids=[rank],
broadcast_buffers=True)
def check_same_model_params():
for pg, ddp_pg in zip(sharded_optimizer.param_groups,
ddp_optimizer.param_groups):
for p, ddp_p in zip(pg["params"], ddp_pg["params"]):
assert torch.allclose(
p, ddp_p, atol=1e-3
), f"Model parameters differ in between DDP and ShardedDDP {p} {ddp_p}"
for b, ddp_b in zip(sharded_ddp_model.buffers(), ddp_model.buffers()):
assert torch.allclose(
b, ddp_b, atol=1e-3
), "Model buffers differ in between DDP and ShardedDDP"
check_same_model_params(
) # The models should stay the same in between the ranks
for i in range(20):
input_tensor = torch.rand((64, 2)).to(device)
# Run DDP
ddp_optimizer.zero_grad()
ddp_optimizer_2.zero_grad()
ddp_loss = ddp_model(input_tensor).abs().sum()
ddp_loss.backward()
ddp_optimizer.step()
ddp_optimizer_2.step()
# Run Sharded
sharded_optimizer.zero_grad()
sharded_optimizer_2.zero_grad()
sharded_loss = sharded_ddp_model(input_tensor).abs().sum()
sharded_loss.backward()
sharded_optimizer.step()
sharded_optimizer_2.step()
check_same_model_params()
dist.destroy_process_group()
from .base import make_test_single_space, make_test_multi_space from pyrl.transforms import OneHot, UnOneHot from torch.nn import Sequential m = lambda space: Sequential(OneHot(space), UnOneHot(space)) test_single_space = make_test_single_space(m) test_multi_space = make_test_multi_space(m)
def __init__(self):
super().__init__()
self.mlp = Sequential(Linear(2, 3), Linear(3, 3), Linear(3, 3),
Linear(3, 3), Linear(3, 3), Linear(3, 3))
def __init__(self, dim):
super(NetGIN, self).__init__()
num_features = 83
nn1_1 = Sequential(Linear(num_features, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn1_2 = Sequential(Linear(num_features, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.conv1_1 = GINConv(nn1_1, train_eps=True)
self.conv1_2 = GINConv(nn1_2, train_eps=True)
self.mlp_1 = Sequential(Linear(2 * dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn2_1 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn2_2 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.conv2_1 = GINConv(nn2_1, train_eps=True)
self.conv2_2 = GINConv(nn2_2, train_eps=True)
self.mlp_2 = Sequential(Linear(2 * dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn3_1 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn3_2 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.conv3_1 = GINConv(nn3_1, train_eps=True)
self.conv3_2 = GINConv(nn3_2, train_eps=True)
self.mlp_3 = Sequential(Linear(2 * dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn4_1 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn4_2 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.conv4_1 = GINConv(nn4_1, train_eps=True)
self.conv4_2 = GINConv(nn4_2, train_eps=True)
self.mlp_4 = Sequential(Linear(2 * dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn5_1 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn5_2 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.conv5_1 = GINConv(nn5_1, train_eps=True)
self.conv5_2 = GINConv(nn5_2, train_eps=True)
self.mlp_5 = Sequential(Linear(2 * dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn6_1 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
nn6_2 = Sequential(Linear(dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.conv6_1 = GINConv(nn6_1, train_eps=True)
self.conv6_2 = GINConv(nn6_2, train_eps=True)
self.mlp_6 = Sequential(Linear(2 * dim, dim), torch.nn.BatchNorm1d(dim), ReLU(), Linear(dim, dim),
torch.nn.BatchNorm1d(dim), ReLU())
self.set2set = Set2Set(1 * dim, processing_steps=6)
self.fc1 = Linear(2 * dim, dim)
self.fc4 = Linear(dim, 12)
print(val_x.shape, val_y.shape)
# loading the pre-trained ResNet18 model
model = models.resnet18(pretrained=True)
# Freeze model weights
for param in model.parameters():
param.requires_grad = False
# checking if GPU is available
if torch.cuda.is_available():
model = model.cuda()
# Add fine-tuning FC layers
model.fc = Sequential(nn.Linear(512, 1000), nn.ReLU(True), nn.Dropout(),
nn.Linear(1000, 4096), nn.ReLU(True), nn.Dropout(),
nn.Linear(4096, 2), nn.LogSoftmax(dim=1))
model.fc = model.fc.cuda()
for param in model.fc.parameters():
param.requires_grad = True
# Check how does our new model look like
print(model)
# batch_size
batch_size = 128
# specify loss function (categorical cross-entropy)
criterion = CrossEntropyLoss()
def __init__(self, edge_dim, dim_init, dim):
super(GINE0Conv, self).__init__(aggr="add")
self.edge_encoder = Sequential(Linear(edge_dim, dim_init), ReLU(), Linear(dim_init, dim_init), ReLU(),
BN(dim_init))
self.mlp = Sequential(Linear(dim_init, dim), ReLU(), Linear(dim, dim), ReLU(), BN(dim))