def test_big(): model = load_t_net(file=True) #print(model.module.state_dict().key()) #print(list(model.module.seq[0].named_parameters())) #print(list(model.module.named_modules())) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") x = torch.randn(1, 3, 384, 224).to(device) model = model.to(device) summary(model, x) num = 0 for name, mod in model.module.named_modules(): num += 1 if num % 40 == 0: continue if isinstance(mod, torch.nn.Conv2d): print("yes") prune.l1_unstructured(mod, name='weight', amount=0.5) else: print("no") print("++++++++++++++++++++++++++++++++++++++++++++++++++++=") parameters_to_prune = ( (model.module.seq[3].list[0][3].list[0][3].list[0][1].convs[2][3], 'weight'), (model.module.seq[3].list[0][3].list[0][3].list[0][1].convs[3][0], 'weight'), (model.module.seq[3].list[0][3].list[1][1].convs[0][0], 'weight'), (model.module.seq[0], 'weight')) #prune.global_unstructured(parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2) x = torch.randn(1, 3, 384, 224).to(device) summary(model, x) #prune.remove(model,'weight') torch.save(model.module, 'gj_dir/after.pth.tar')
def main(): device = torch.device("cuda") # Dataset test_transform = transforms.Compose([transforms.ToTensor()]) test_dataset = PHD08(options["test_dir"], test_transform) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=options["batch_size"], shuffle=False, num_workers=options["workers"], pin_memory=True, ) checkpoint_path = "%s.pth" % options["name"] # Model if options["mode"] == "Normal": model = Normal(options["hidden_dim"], options["context_dim"]) else: model = VAE( options["hidden_dim"], options["context_dim"], options["addition_dim"] ) model.load_state_dict(torch.load(checkpoint_path)) model.eval() summary(model, torch.zeros((1, 1, 28, 28))) model = model.to(device) # Train test_model(options["mode"], device, model, test_loader)
def __init__(self, module, opt): self.opt = opt self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.net = module.Net(opt).to(self.dev) summary(self.net, torch.zeros(4, 3, 64, 64).to(self.dev)) if opt.pretrain: self.load(opt.pretrain) self.loss_fn = nn.L1Loss() self.optim = torch.optim.Adam(self.net.parameters(), opt.lr, betas=(0.9, 0.999), eps=1e-8) self.scheduler = torch.optim.lr_scheduler.MultiStepLR( self.optim, [1000 * int(d) for d in opt.decay.split("-")], gamma=opt.gamma, ) if not opt.test_only: self.train_loader = generate_loader("train", opt) self.test_loader = generate_loader("test", opt) self.t1, self.t2 = None, None self.best_psnr, self.best_step = 0, 0
def test(self,device='cpu'): input_tensor = torch.rand(1, self.input_channels, 12, 12, 12) ideal_out = torch.rand(1, self.num_classes, 12, 12, 12) out = self.forward(input_tensor) assert ideal_out.shape == out.shape summary(self.to(torch.device(device)), (self.input_channels, 12, 12, 12),device=device) import torchsummaryX torchsummaryX.summary(self, input_tensor.to(device)) print("DenseNet3D-2 test is complete!!!!\n\n\n\n\n")
def main(): args = parser.parse_args() # create model model = resnet50(num_classes=1000) model.eval() summary(model, torch.zeros((1, 3, 224, 224))) model_weight = torch.load('ResNet50-CURL-1G.pth') model = torch.nn.DataParallel(model.cuda(), device_ids=range(torch.cuda.device_count())) model.load_state_dict(model_weight) cudnn.benchmark = True # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() # Data loading code from lmdb data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), # transforms.Resize(256), # transforms.RandomCrop((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } data_dir = args.data_base print("| Preparing model...") dsets = { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val'] } train_loader = torch.utils.data.DataLoader(dsets['train'], batch_size=args.batch_size, shuffle=True, num_workers=8, pin_memory=True) val_loader = torch.utils.data.DataLoader(dsets['val'], batch_size=args.batch_size, shuffle=False, num_workers=8, pin_memory=True) print('data_loader_success!') # evaluate and train validate(val_loader, model, criterion)
def test(self, device='cpu'): input_tensor = torch.rand(1, 2, 32, 32, 32) ideal_out = torch.rand(1, self.num_classes, 32, 32, 32) out = self.forward(input_tensor) assert ideal_out.shape == out.shape summary(self.to(torch.device(device)), (2, 32, 32, 32), device=device) import torchsummaryX torchsummaryX.summary(self, input_tensor.to(device)) print("SkipDenseNet3D test is complete")
def test(self, device='cpu'): a = torch.rand(1, self.in_channels, 8, 8, 8) ideal_out = torch.rand(1, self.classes, 8, 8, 8) summary(self.to(torch.device(device)), (self.in_channels, 8, 8, 8), device=device) b, c = self.forward(a) import torchsummaryX torchsummaryX.summary(self, a.to(device)) assert ideal_out.shape == b.shape assert ideal_out.shape == c.shape print("Test DenseVoxelNet is complete")
def test_args(): class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(64, 64, 3, 1, 1) def forward(self, x, args1, args2): out = self.conv1(x) out = self.conv1(out) return out summary(Net(), torch.zeros((1, 64, 28, 28)), "args1", args2="args2")
def __init__(self, module, config): self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") kwargs = { "num_channels": config.num_channels, "groups": config.group, "mobile": config.mobile } if config.scale > 0: kwargs["scale"] = config.scale else: kwargs["multi_scale"] = True self.G = module.Net(**kwargs).to(self.device) if config.msd: self.D = [ module.Discriminator(1).to(self.device), module.Discriminator(2).to(self.device), module.Discriminator(4).to(self.device), ] else: self.D = [module.Discriminator().to(self.device)] load(self.G, config.pretrained_ckpt) self.loss_l2 = nn.MSELoss() self.loss_gan = utils.GANLoss(self.device) self.loss_vgg = utils.VGGLoss().to(self.device) self.optim_G = torch.optim.Adam(self.G.parameters(), config.lr) D_params = list() for D in self.D: D_params += list(D.parameters()) self.optim_D = torch.optim.Adam(D_params, config.lr) self.train_loader = generate_loader(path=config.train_data, scale=config.scale, train=True, size=config.patch_size, batch_size=config.batch_size, num_workers=1, shuffle=True, drop_last=True) self.step = 0 self.config = config summary(self.G, torch.zeros((1, 3, 720 // 4, 1280 // 4)).to(self.device), scale=4) self.writer = SummaryWriter( log_dir=os.path.join("./runs", config.memo)) os.makedirs(config.ckpt_dir, exist_ok=True)
def compute(self, input_shape, index=[-1]): model_handle = self._model.register_forward_pre_hook( self._add_pre_forward(self._model)) replace_names_dict(self._model) # 注册模型的hook self._model.apply(self._add_model_forward(get_names_dict(self._model))) extra = torch.zeros( (2, input_shape[2], input_shape[0], input_shape[1])) with torch.no_grad(): self._model(extra) # 删除hook,防止影响 for handle in self._hooks: handle.remove() if summaryX is None: raise ModuleNotFoundError('Please install torchsummaryX,' ' pip install torchsummaryX ' '') summaryX.summary(self._model, extra) # 添加hook if not isinstance(index, list): index = [index] handles = [] print('--请不要查看FC层的感受野,因为其感受野永远是全图大小--') print('你查看的特征图感受野对应的名称为:') for i in index: if i < 0: i = len(list(self._summary.keys())) + i print(list(self._summary.keys())[i]) m = self._summary.get(list(self._summary.keys())[i]) handles.append( m.register_forward_hook(self._add_layer_forward(self._model))) def get_model_fn(): self._model.eval() return self._model rf = PytorchReceptiveField(get_model_fn) rf_params = rf.compute(input_shape=input_shape) print('对应的感受野如下所示(看每一行最后一个Size即可):') for maps_desc in rf.feature_maps_desc: size = maps_desc[1][2] if size[0] >= input_shape[0] and size[1] >= input_shape[1]: print('感受野计算可能不正确,请扩大输入图片shape,当前是{},{}'.format( input_shape[0], input_shape[1])) print(maps_desc) # 移除hook for handle in handles: handle.remove() model_handle.remove() return rf_params
def extract_logits(): global args, best_prec1 args = parser.parse_args() args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size) # create model model_weight = torch.load(FLAGS['pretrain_model']) # 78.771% model = resnet50(model_weight, num_classes=FLAGS['class_num']).cuda() model.eval() summary(model, torch.zeros((1, 3, 224, 224)).cuda()) # test_accuracy(model) cudnn.benchmark = True # Data loading code from lmdb data_list = np.load( os.path.join(FLAGS['6x_larger_dataset'], 'data_list.npy')) np.save('data_list.npy', data_list) data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), # transforms.Resize(256), # transforms.RandomCrop((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } print("| Preparing model...") dsets = {} dsets['val'] = MyDataSet(data_transforms['val']) val_loader = torch.utils.data.DataLoader(dsets['val'], batch_size=4 * FLAGS['batch_size'], shuffle=False, num_workers=8, pin_memory=True) print('data_loader_success!') validate(val_loader, model)
def main(): global args, best_prec1 # Phase 1 : Model setup print('\n[Phase 2] : Model setup') model = mobilenetv2.MobileNetV2(args.ft_model_path) model.eval() summary(model, torch.zeros((1, 3, 224, 224))) model_ft = torch.nn.DataParallel(model.cuda()) cudnn.benchmark = True print("model setup success!") # Phase 2 : Data Load # Data pre-processing data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), # transforms.Resize(256), # transforms.RandomCrop((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } data_dir = args.data_base print("| Preparing data...") dsets = { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val'] } val_loader = torch.utils.data.DataLoader(dsets['val'], batch_size=args.batch_size, shuffle=False, num_workers=8, pin_memory=True) print('data_loader_success!') # Phase 3: fine_tune model print('\n[Phase 3] : Model fine tune') # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() validate(val_loader, model_ft, criterion)
def test(self, device='cpu'): a = torch.rand(1, self.in_channels, 8, 8, 8) ideal_out = torch.rand(1, self.classes, 8, 8, 8) summary(self.to(torch.device(device)), (self.in_channels, 8, 8, 8), device=device) # TODO put Auxiliary mid-layer prediction aside temp # b, c = self.forward(a) b = self.forward(a) import torchsummaryX torchsummaryX.summary(self, a.to(device)) assert ideal_out.shape == b.shape # TODO put Auxiliary mid-layer prediction aside temp # assert ideal_out.shape == c.shape print("Test DenseVoxelNet is complete")
def test_BertForUtteranceLanguageModeling(): num_samples = 5 seq_len = 64 config = BertConfig(max_position_embeddings=512) model = BertForUtteranceLanguageModeling(config) input_ids = np.ones(shape=(num_samples, seq_len), dtype=np.int32) attention_masks = np.ones(shape=(num_samples, seq_len), dtype=np.int32) label_ids = np.full(shape=(num_samples, ), dtype=np.int32, fill_value=0) model.cpu() model.float() summary(model, torch.tensor(input_ids.astype(np.int64)), torch.tensor(attention_masks.astype(np.int64)), torch.tensor(label_ids.astype(np.int64)))
def _print_model_structure(self, print_summary=True): import torchsummaryX as summaryX self._net.apply(self._add_model_forward(get_names_dict(self._net))) extra = torch.zeros( (2, self._num_channel, self._input_shape[0], self._input_shape[1])) if self._use_gpu: extra = extra.cuda() with torch.no_grad(): self._net(extra) # 删除hook,防止影响 for handle in self._hooks: handle.remove() if print_summary: summaryX.summary(self._net, extra)
def main(): args = parse_args() update_config(cfg, args) cfg.defrost() cfg.freeze() model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(cfg, is_train=False) model.eval() dump_input = torch.rand( (1, 3, cfg.DATASET.INPUT_SIZE, cfg.DATASET.INPUT_SIZE)) summary(model, dump_input)
def test_lstm(): class Net(nn.Module): def __init__(self, vocab_size=20, embed_dim=300, hidden_dim=512, num_layers=2): super().__init__() self.hidden_dim = hidden_dim self.embedding = nn.Embedding(vocab_size, embed_dim) self.encoder = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers) self.decoder = nn.Linear(hidden_dim, vocab_size) def forward(self, x): embed = self.embedding(x) out, hidden = self.encoder(embed) out = self.decoder(out) out = out.view(-1, out.size(2)) return out, hidden inputs = torch.zeros((100, 1), dtype=torch.long) # [length, batch_size] df, df_total = summary(Net(), inputs) assert df.shape[0] == 3, 'Should find 3 layers'
def print_net_summary(log, net, x: Dict[str, Tensor]): """ Print the net summary into a log file Args: log: (str) log file path net: (nn.Module) network instance x: (dict of torch.Tensor) input tensor of size [batch, seq_length, hidden_dim] """ original_stdout = sys.stdout inputs = x['input_ids'].to('cuda') with open(log, 'w') as f: sys.stdout = f summary(net, inputs) sys.stdout = original_stdout f.close()
def main(args): traj, n_traj, translator = traj_prepare( TRAJ_PATH, CODE_PATH, NUM_SAMPLES, use_cols=args.begin_index + np.arange(args.num_per_day) * args.time_interval, add_idx=ADD_IDX) NUM_CLASS = n_traj num_training_sample = int(NUM_SAMPLES * 0.8) ix_sos = translator.trans_code2ix(["0/SOS"]) traj = np.concatenate([np.ones((traj.shape[0], 1), dtype=int), traj], axis=1) training_data, validation_data = traj[:num_training_sample, :-1], traj[ num_training_sample:, :-1] training_label, validation_label = traj[:num_training_sample, 1:], traj[num_training_sample:, 1:] training_loader = torch.utils.data.DataLoader(TrajDatasetGeneral( [training_data, training_label]), batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=False) validation_loader = torch.utils.data.DataLoader(TrajDatasetGeneral( [validation_data, validation_label]), batch_size=2000, shuffle=True, num_workers=4, drop_last=False) seq_vae = SeqVAE(num_class=NUM_CLASS, embedding_size=args.embed_size, embedding_drop_out=args.emb_dropout, hid_size=args.hid_size, hid_layers=args.hid_layer, latent_z_size=args.latent_z_size, word_dropout_rate=args.word_dropout, gru_drop_out=args.gru_dropout, anneal_k=args.anneal_k, anneal_x0=args.anneal_x0, anneal_function=args.anneal_func, data_keys=DATA_KEYS, learning_rate=args.lr, writer_comment="SentenceVAE", sos_idx=ix_sos, unk_idx=None) test_batch = wrap_data(next(iter(training_loader)), DATA_KEYS) model_summary = torchsummaryX.summary(seq_vae.model, test_batch['data']) seq_vae.model.to(DEVICE) seq_vae.fit(training_loader=training_loader, validation_loader=validation_loader, epochs=args.epochs, device=DEVICE) return
def __init__(self, module, config): self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) kwargs = { "num_channels": config.num_channels, "groups": config.group, "mobile": config.mobile } if config.scale > 0: kwargs["scale"] = config.scale else: kwargs["multi_scale"] = True self.net = module.Net(**kwargs).to(self.device) init_weights(self.net, config.init_type, config.init_scale) self.loss_fn = nn.L1Loss() self.optim = torch.optim.Adam( filter(lambda p: p.requires_grad, self.net.parameters()), config.lr ) self.scheduler = torch.optim.lr_scheduler.StepLR( self.optim, config.decay, gamma=0.5, ) self.train_loader = generate_loader( path=config.train_data, scale=config.scale, train=True, size=config.patch_size, batch_size=config.batch_size, num_workers=1, shuffle=True, drop_last=True ) self.step = 0 self.config = config summary( self.net, torch.zeros((1, 3, 720//4, 1280//4)).to(self.device), scale=4 ) self.writer = SummaryWriter(log_dir=os.path.join("./runs", config.memo)) os.makedirs(config.ckpt_dir, exist_ok=True)
def test_resnet(): model = torchvision.models.resnet50() df, df_total = summary(model, torch.zeros(4, 3, 224, 224)) # According to https://arxiv.org/abs/1605.07146, resnet50 has ~25.6 M trainable params. # Lets make sure we count them correctly np.testing.assert_approx_equal(25.6e6, df_total['Totals']['Total params'], significant=3)
def init_model(self): criterion = nn.CrossEntropyLoss().to(self.device) if self.args.stage == "search": model = Network(self.args.init_channels, 10, self.args.layers, criterion, self.device) else: genotype = genotypes.FedNAS_V1 logging.info(genotype) model = NetworkCIFAR(self.args.init_channels, 10, self.args.layers, self.args.auxiliary, genotype) model.to(self.device) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") summary(model, torch.zeros(2, 3, 32, 32).to(self.device)) # summary(model, (3, 224, 224), batch_size=64) return model
def test_convnet(): class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) summary(Net(), torch.zeros((1, 1, 28, 28)))
def prune_global(): model = load_t_net(file=True) pr_list = [] for name, mod in model.module.named_modules(): print(name, mod) pr_tup = (mod, "weight") pr_list.append(pr_tup) parameters_to_prune = ( (model.module.seq[3].list[0][3].list[0][3].list[0][1].convs[2][3], 'weight'), (model.module.seq[3].list[0][3].list[0][3].list[0][1].convs[3][0], 'weight'), (model.module.seq[3].list[0][3].list[1][1].convs[0][0], 'weight'), (model.module.seq[0], 'weight')) prune.global_unstructured(pr_list, pruning_method=prune.L1Unstructured, amount=0.5) x = torch.randn(1, 3, 384, 224).cuda() summary(model, x) # prune.remove(model,'weight') torch.save(model.module, 'gj_dir/after_glo.pth.tar')
def new_prune(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() # 1 input image channel, 6 output channels, 3x3 square conv kernel self.conv1 = nn.Conv2d(1, 6, 3) self.conv2 = nn.Conv2d(6, 16, 3) self.fc1 = nn.Linear(16 * 5 * 5, 120) # 5x5 image dimension self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) x = F.max_pool2d(F.relu(self.conv2(x)), 2) x = x.view(-1, int(x.nelement() / x.shape[0])) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x model = LeNet().to(device=device) x = torch.randn(1, 1, 28, 28).to(device) summary(model, x) parameters_to_prune = ( (model.conv1, 'weight'), (model.conv2, 'weight'), (model.fc1, 'weight'), (model.fc2, 'weight'), (model.fc3, 'weight'), ) prune.global_unstructured( parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2, ) x = torch.randn(1, 1, 28, 28).to(device) summary(model, x)
def prune_6M(para): model = load_t_net(file=True) x = torch.randn(1, 3, 384, 224).cuda() model = model.cuda() summary(model, x) num = 0 for name, mod in model.module.named_modules(): num += 1 if num % para == 0: continue if isinstance(mod, torch.nn.Conv2d): print("yes") prune.l1_unstructured(mod, name='weight', amount=0.5) else: print("no") print("++++/++++/+++++/+++++/++++++++++++=+++++++++++++++++++++=") x = torch.randn(1, 3, 384, 224).cuda() summary(model, x) # prune.remove(model,'weight') torch.save(model.module, 'gj_dir/after_6M.pth.tar')
def test_BertForUtteranceOrderPrediction(): num_samples = 5 num_utterances = 6 seq_len = 64 max_utterances = 10 config = BertConfig(max_position_embeddings=512) utterance_config = BertConfig(max_position_embeddings=max_utterances + 1, num_hidden_layers=2) model = BertForUtteranceOrderPrediction(config, utterance_config, max_utterances) utterances_input_ids = np.ones(shape=(num_samples, num_utterances, seq_len), dtype=np.int32) attention_masks = np.ones(shape=(num_samples, num_utterances, seq_len), dtype=np.int32) label_ids = np.full(shape=(num_samples, ), dtype=np.int32, fill_value=0) model.cpu() model.float() summary(model, torch.tensor(utterances_input_ids.astype(np.int64)), torch.tensor(attention_masks.astype(np.int64)), torch.tensor(label_ids.astype(np.int64)))
def test_get_two_stream(): net = get_two_stream() bridge = net.bridge_frame2op print("net: ", bridge) for param in bridge.parameters(): print(type(param.data), param.size()) # bridge = net.bridge_op2frame print("net: ", bridge) for param in bridge.parameters(): print(type(param.data), param.size()) # # 预估参数量 from torchsummaryX import summary x = torch.randn(2, 64, 64, 64) summary(net.bridge_frame2op, x) print("ok \n\n") # summary(net.bridge_op2frame, x) print("ok \n\n") # b, t, c, h, w = 2, 9, 3, 256, 256 frame = torch.randn(b, t, c, h, w) b, t, c, h, w = 2, 8, 2, 256, 256 op = torch.randn(b, t, c, h, w) summary(net, [frame, op]) # 测试 net的 forward # out_list = net([frame, op]) for out in out_list: print(out.size())
def test_get_two_stream_v1(): net = get_two_stream() bridge = net.bridge_frame2op print("net: ", bridge) # top, bottom = bridge[0], bridge[1] # print("top: ", top) # print("bottom: ", bottom) for param in bridge.parameters(): print(type(param.data), param.size()) # bridge = net.bridge_op2frame print("net: ", bridge) # top, bottom = bridge[0], bridge[1] # print("top: ", top) # print("bottom: ", bottom) for param in bridge.parameters(): print(type(param.data), param.size()) # # 预估参数量 from torchsummaryX import summary x = torch.randn(1, 1, 32, 32) # [b,1,h,w] y = torch.randn(1, 1, 64, 64) summary(net.bridge_frame2op, [x, y]) print("ok \n\n") # summary(net.bridge_op2frame, [x, y]) print("ok \n\n") # x = torch.randn(1, 3, 256, 256) y = torch.randn(1, 2, 256, 256) # 当前是 2-channel summary(net, x, y) # 测试 net的 forward print("ok \n\n")
def inference(net, config): device = torch.device("cpu") net = net.to(device) inp = torch.zeros( (1, 3, config.height//config.scale, config.width//config.scale) ).to(device) summary(net, inp, scale=4) with torch.no_grad(): net(inp, config.scale) avg_time = 0.0 for _ in range(config.num_exp): t1 = time.time() with torch.no_grad(): net(inp, config.scale) t2 = time.time() avg_time += (t2-t1)/config.num_exp print("{} x{}: {:.3f}".format(config.model, config.scale, avg_time))