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'] = []