def __init__(self, n_input, D, n_hidden, n_layers): super().__init__() self.D = D #latent space dimension self.encoder = MLP(n_input, D * 2, n_hidden, n_layers) # output is mean and std of z self.decoder = MLP(D, n_input, n_hidden, n_layers) self.sigmoid = nn.Sigmoid() self.device = torch.device('cuda:' + str(get_gpu_memory_map()) if torch.cuda.is_available() else 'cpu') self.to(self.device)
def __init__(self, n_channels, n_layers, kernel_size, bias=False): assert n_layers > 0, 'The net must have at least one layer!' super().__init__() self.device = torch.device( 'cuda:' + str(get_gpu_memory_map()) if torch.cuda.is_available() else 'cpu') self.activation = nn.ReLU() self.sigmoid = nn.Sigmoid() self.initialize_net(n_channels, kernel_size, n_layers, bias) self.to(self.device)
help='number of update steps for checkpoint and validation' ) parser.add_argument('--patience', type=int, action='store', default=5, help='patience for early stopping') parser.add_argument('--debug', action='store_true', help='run in debugging mode') parser.add_argument('--no-cuda', action='store_true', help='disable GPU training') args = parser.parse_args() args.bi = not args.no_bi args.cuda = torch.cuda.is_available() and not args.no_cuda if args.cuda: os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" sorted_gpu_info = get_gpu_memory_map() for gpu_id, (mem_left, util) in sorted_gpu_info: if mem_left >= 1000: os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) # args.cuda = f'cuda:{gpu_id}' print('use gpu:{} with {} MB left, util {}%'.format(gpu_id, mem_left, util)) break else: print(f'no gpu has memory left >= {args.min_mem} MB, exiting...') exit() else: torch.set_num_threads(10) args.data = 'train010203x_test04' args.remark = 'del' remark = f'_{args.data}_{args.remark}'
parser = argparse.ArgumentParser() parser.add_argument('--GPU', type=str, default='gpu_name') parser.add_argument('--mode', type=str, default='FP32', choices=['FP32', 'FP16', 'amp']) parser.add_argument('--batch_size', type=int, default=128) parser.add_argument('--iteration', type=int, default=100) parser.add_argument('--opt_level', type=str, default="O1", help='What type of half precision to use.') args = parser.parse_args() print('------------ Options -------------') for k, v in sorted(vars(args).items()): print('%s: %s' % (str(k), str(v))) print('-------------- End ---------------') init_mem = get_gpu_memory_map() seed = 7 np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.benchmark = True transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) train_loader = DataLoader(cifar.CIFAR10(root='cifar', train=True, transform=transform, download=True), batch_size=args.batch_size, shuffle=False, pin_memory=True, drop_last=True) test_loader = DataLoader(cifar.CIFAR10(root='cifar', train=False, transform=transform, download=True), batch_size=args.batch_size, shuffle=False, pin_memory=True, drop_last=True) loss_fn = torch.nn.CrossEntropyLoss().cuda() result = {} result['train_time'] = []