Download lib/network.py from SJTU-TES/TMM: direct link, hf CLI and curl.
- Browser
- Download file 13 kB
-
https://huggingface.co/spaces/SJTU-TES/TMM/resolve/main/lib/network.py
- Command line
-
hf download hf://spaces/SJTU-TES/TMM/lib/network.py
-
curl -L -o network.py https://huggingface.co/spaces/SJTU-TES/TMM/resolve/main/lib/network.py
13 kB
| import sys | |
| thismodule = sys.modules[__name__] | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import random | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| torch.manual_seed(123) | |
| def get_autoencoder(config): | |
| ae_cls = getattr(thismodule, config.autoencoder.cls) | |
| return ae_cls(config.autoencoder) | |
| class ConvEncoder(nn.Module): | |
| def build_from_config(cls, config): | |
| conv_pool = None if config.conv_pool is None else getattr(nn, config.conv_pool) | |
| encoder = cls(config.channels, config.padding, config.kernel_size, config.conv_stride, conv_pool) | |
| return encoder | |
| def __init__(self, channels, padding=3, kernel_size=8, conv_stride=2, conv_pool=None): | |
| super(ConvEncoder, self).__init__() | |
| self.in_channels = channels[0] | |
| model = [] | |
| acti = nn.LeakyReLU(0.2) | |
| nr_layer = len(channels) - 1 | |
| for i in range(nr_layer): | |
| if conv_pool is None: | |
| model.append(nn.ReflectionPad1d(padding)) | |
| model.append(nn.Conv1d(channels[i], channels[i+1], kernel_size=kernel_size, stride=conv_stride)) | |
| model.append(acti) | |
| else: | |
| model.append(nn.ReflectionPad1d(padding)) | |
| model.append(nn.Conv1d(channels[i], channels[i+1], kernel_size=kernel_size, stride=conv_stride)) | |
| model.append(acti) | |
| model.append(conv_pool(kernel_size=2, stride=2)) | |
| self.model = nn.Sequential(*model) | |
| def forward(self, x): | |
| x = x[:, :self.in_channels, :] | |
| x = self.model(x) | |
| return x | |
| class ConvDecoder(nn.Module): | |
| def build_from_config(cls, config): | |
| decoder = cls(config.channels, config.kernel_size) | |
| return decoder | |
| def __init__(self, channels, kernel_size=7): | |
| super(ConvDecoder, self).__init__() | |
| model = [] | |
| pad = (kernel_size - 1) // 2 | |
| acti = nn.LeakyReLU(0.2) | |
| for i in range(len(channels) - 1): | |
| model.append(nn.Upsample(scale_factor=2, mode='nearest')) | |
| model.append(nn.ReflectionPad1d(pad)) | |
| model.append(nn.Conv1d(channels[i], channels[i + 1], | |
| kernel_size=kernel_size, stride=1)) | |
| if i == 0 or i == 1: | |
| model.append(nn.Dropout(p=0.2)) | |
| if not i == len(channels) - 2: | |
| model.append(acti) # whether to add tanh a last? | |
| #model.append(nn.Dropout(p=0.2)) | |
| self.model = nn.Sequential(*model) | |
| def forward(self, x): | |
| return self.model(x) | |
| class Discriminator(nn.Module): | |
| def __init__(self, config): | |
| super(Discriminator, self).__init__() | |
| self.gan_type = config.gan_type | |
| encoder_cls = getattr(thismodule, config.encoder_cls) | |
| self.encoder = encoder_cls.build_from_config(config) | |
| self.linear = nn.Linear(config.channels[-1], 1) | |
| def forward(self, seqs): | |
| code_seq = self.encoder(seqs) | |
| logits = self.linear(code_seq.permute(0, 2, 1)) | |
| return logits | |
| def calc_dis_loss(self, x_gen, x_real): | |
| fake_logits = self.forward(x_gen) | |
| real_logits = self.forward(x_real) | |
| if self.gan_type == 'lsgan': | |
| loss = torch.mean((fake_logits - 0) ** 2) + torch.mean((real_logits - 1) ** 2) | |
| elif self.gan_type == 'nsgan': | |
| all0 = torch.zeros_like(fake_logits, requires_grad=False) | |
| all1 = torch.ones_like(real_logits, requires_grad=False) | |
| loss = torch.mean(F.binary_cross_entropy(F.sigmoid(fake_logits), all0) + | |
| F.binary_cross_entropy(F.sigmoid(real_logits), all1)) | |
| else: | |
| raise NotImplementedError | |
| return loss | |
| def calc_gen_loss(self, x_gen): | |
| logits = self.forward(x_gen) | |
| if self.gan_type == 'lsgan': | |
| loss = torch.mean((logits - 1) ** 2) | |
| elif self.gan_type == 'nsgan': | |
| all1 = torch.ones_like(logits, requires_grad=False) | |
| loss = torch.mean(F.binary_cross_entropy(F.sigmoid(logits), all1)) | |
| else: | |
| raise NotImplementedError | |
| return loss | |
| class Autoencoder3f(nn.Module): | |
| def __init__(self, config): | |
| super(Autoencoder3f, self).__init__() | |
| assert config.motion_encoder.channels[-1] + config.body_encoder.channels[-1] + \ | |
| config.view_encoder.channels[-1] == config.decoder.channels[0] | |
| self.n_joints = config.decoder.channels[-1] // 3 | |
| self.body_reference = config.body_reference | |
| motion_cls = getattr(thismodule, config.motion_encoder.cls) | |
| body_cls = getattr(thismodule, config.body_encoder.cls) | |
| view_cls = getattr(thismodule, config.view_encoder.cls) | |
| self.motion_encoder = motion_cls.build_from_config(config.motion_encoder) | |
| self.body_encoder = body_cls.build_from_config(config.body_encoder) | |
| self.view_encoder = view_cls.build_from_config(config.view_encoder) | |
| self.decoder = ConvDecoder.build_from_config(config.decoder) | |
| self.body_pool = getattr(F, config.body_encoder.global_pool) if config.body_encoder.global_pool is not None else None | |
| self.view_pool = getattr(F, config.view_encoder.global_pool) if config.view_encoder.global_pool is not None else None | |
| def forward(self, seqs): | |
| return self.reconstruct(seqs) | |
| def encode_motion(self, seqs): | |
| motion_code_seq = self.motion_encoder(seqs) | |
| return motion_code_seq | |
| def encode_body(self, seqs): | |
| body_code_seq = self.body_encoder(seqs) | |
| kernel_size = body_code_seq.size(-1) | |
| body_code = self.body_pool(body_code_seq, kernel_size) if self.body_pool is not None else body_code_seq | |
| return body_code, body_code_seq | |
| def encode_view(self, seqs): | |
| view_code_seq = self.view_encoder(seqs) | |
| kernel_size = view_code_seq.size(-1) | |
| view_code = self.view_pool(view_code_seq, kernel_size) if self.view_pool is not None else view_code_seq | |
| return view_code, view_code_seq | |
| def decode(self, motion_code, body_code, view_code): | |
| if body_code.size(-1) == 1: | |
| body_code = body_code.repeat(1, 1, motion_code.shape[-1]) | |
| if view_code.size(-1) == 1: | |
| view_code = view_code.repeat(1, 1, motion_code.shape[-1]) | |
| complete_code = torch.cat([motion_code, body_code, view_code], dim=1) | |
| out = self.decoder(complete_code) | |
| return out | |
| def cross3d(self, x_a, x_b, x_c): | |
| motion_a = self.encode_motion(x_a) | |
| body_b, _ = self.encode_body(x_b) | |
| view_c, _ = self.encode_view(x_c) | |
| out = self.decode(motion_a, body_b, view_c) | |
| return out | |
| def cross2d(self, x_a, x_b, x_c): | |
| motion_a = self.encode_motion(x_a) | |
| body_b, _ = self.encode_body(x_b) | |
| view_c, _ = self.encode_view(x_c) | |
| out = self.decode(motion_a, body_b, view_c) | |
| batch_size, channels, seq_len = out.size() | |
| n_joints = channels // 3 | |
| out = out.view(batch_size, n_joints, 3, seq_len) | |
| out = out[:, :, [0, 2], :] | |
| out = out.view(batch_size, n_joints * 2, seq_len) | |
| return out | |
| def cross2d_adv(self, x_a, x_b, x_c): | |
| x_a.cpu() | |
| x_a_shape = x_a.shape | |
| print(x_a.shape) | |
| #motion_a_org = self.encode_motion(x_a) | |
| print(x_a) | |
| # The heatmap image is saved as 'tensor_heatmap.png' in the current directory | |
| # for i in range(0,119): | |
| # x_a[0][11][i]+=1 | |
| #x_a[0][7][60]+=0.01 | |
| #motion_a = self.encode_motion(x_a) | |
| # print(motion_a.shape) | |
| # print(motion_a[0][0]-motion_a_org[0][0]) | |
| # res = motion_a[0] - motion_a_org[0] | |
| # res = res.cpu().detach().numpy() | |
| # # Code for plotting the heatmap | |
| # plt.figure(figsize=(15, 10)) | |
| # plt.imshow(res, cmap='hot', interpolation='nearest') | |
| # plt.colorbar() | |
| # plt.title('Heatmap of the Tensor') | |
| # # Save the heatmap to a local file | |
| # plt.savefig('/home/fazhong/studio/transmomo.pytorch/tensor_heatmap2.png') | |
| # plt.close() | |
| initial_motion_a = self.encode_motion(x_a) # 计算初始的motion_a | |
| # 定义一个函数来计算motion的变化量 | |
| def motion_change(motion_a, initial_motion_a): | |
| return (motion_a - initial_motion_a).norm() | |
| # 设置初始的最大变化量为0 | |
| max_change = 0 | |
| # 扰动次数,可以根据需要更改 | |
| num_perturbations = 10000 | |
| init_a = x_a.clone() | |
| for _ in range(num_perturbations): | |
| # 复制x_a以避免在原始数据上修改 | |
| x_a_perturbed = x_a.clone().cpu() | |
| # 选择要扰动的随机点 | |
| batch_idx, seq_idx, feature_idx = (torch.randint(0, x_a.size(0), (1,)), | |
| torch.randint(0, x_a.size(1), (1,)), | |
| torch.randint(0, x_a.size(2), (1,))) | |
| # 在选定点上加上扰动 | |
| x_a_perturbed[batch_idx, seq_idx, feature_idx] += 10 * torch.randn(1) | |
| # 计算扰动后的motion_a | |
| perturbed_motion_a = self.encode_motion(x_a_perturbed.to('cuda:0')) | |
| # 计算变化量 | |
| change = motion_change(perturbed_motion_a, initial_motion_a) | |
| # 如果变化量大于之前保存的最大变化量,则更新x_a和最大变化量 | |
| if change > max_change: | |
| x_a = x_a_perturbed | |
| max_change = change | |
| # 最后,x_a将是导致最大motion_a变化的扰动版本 | |
| # max_change是这个变化量 | |
| # print(max_change) | |
| # print(max_change.shape) | |
| print(x_a_perturbed - init_a.cpu()) | |
| motion_a = self.encode_motion(x_a_perturbed.to('cuda:0')) | |
| # motion_a = self.encode_motion(x_a.to('cuda:0')) | |
| body_b, _ = self.encode_body(x_b) | |
| view_c, _ = self.encode_view(x_c) | |
| out = self.decode(motion_a, body_b, view_c) | |
| batch_size, channels, seq_len = out.size() | |
| n_joints = channels // 3 | |
| out = out.view(batch_size, n_joints, 3, seq_len) | |
| out = out[:, :, [0, 2], :] | |
| out = out.view(batch_size, n_joints * 2, seq_len) | |
| return out | |
| def cross2d_one(self, x_a): | |
| motion_a = self.encode_motion(x_a) | |
| body_b, _ = self.encode_body(x_a) | |
| view_c, _ = self.encode_view(x_a) | |
| out = self.decode(motion_a, body_b, view_c) | |
| batch_size, channels, seq_len = out.size() | |
| n_joints = channels // 3 | |
| out = out.view(batch_size, n_joints, 3, seq_len) | |
| out = out[:, :, [0, 2], :] | |
| out = out.view(batch_size, n_joints * 2, seq_len) | |
| return out | |
| def adv_cross(self,x_a): | |
| motion_a = self.encode_motion(x_a) | |
| body_b, _ = self.encode_body(x_a) | |
| view_c, _ = self.encode_view(x_a) | |
| return motion_a | |
| def reconstruct3d(self, x): | |
| motion_code = self.encode_motion(x) | |
| body_code, _ = self.encode_body(x) | |
| view_code, _ = self.encode_view(x) | |
| out = self.decode(motion_code, body_code, view_code) | |
| return out | |
| def reconstruct2d(self, x): | |
| motion_code = self.encode_motion(x) | |
| body_code, _ = self.encode_body(x) | |
| view_code, _ = self.encode_view(x) | |
| out = self.decode(motion_code, body_code, view_code) | |
| batch_size, channels, seq_len = out.size() | |
| n_joints = channels // 3 | |
| out = out.view(batch_size, n_joints, 3, seq_len) | |
| out = out[:, :, [0, 2], :] | |
| out = out.view(batch_size, n_joints * 2, seq_len) | |
| return out | |
| def interpolate(self, x_a, x_b, N): | |
| step_size = 1. / (N-1) | |
| batch_size, _, seq_len = x_a.size() | |
| motion_a = self.encode_motion(x_a) | |
| body_a, body_a_seq = self.encode_body(x_a) | |
| view_a, view_a_seq = self.encode_view(x_a) | |
| motion_b = self.encode_motion(x_b) | |
| body_b, body_b_seq = self.encode_body(x_b) | |
| view_b, view_b_seq = self.encode_view(x_b) | |
| batch_out = torch.zeros([batch_size, N, N, 2 * self.n_joints, seq_len]) | |
| for i in range(N): | |
| motion_weight = i * step_size | |
| for j in range(N): | |
| body_weight = j * step_size | |
| motion = (1. - motion_weight) * motion_a + motion_weight * motion_b | |
| body = (1. - body_weight) * body_a + body_weight * body_b | |
| view = (1. - body_weight) * view_a + body_weight * view_b | |
| out = self.decode(motion, body, view) | |
| batch_size, channels, seq_len = out.size() | |
| n_joints = channels // 3 | |
| out = out.view(batch_size, n_joints, 3, seq_len) | |
| out = out[:, :, [0, 2], :] | |
| out = out.view(batch_size, n_joints * 2, seq_len) | |
| batch_out[:, i, j, :, :] = out | |
| return batch_out | |