def init_environment(self):
# gpu
if self.arg.use_gpu:
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
else:
self.dev = "cpu"
# random seed
if self.arg.fix_random:
SEED = 0
#torch.backends.cudnn.deterministic = True
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(SEED)
random.seed(SEED)
# dir
self.work_dir = get_dir(self.arg.work_dir, None)
self.config_dir = get_dir(self.arg.work_dir, 'config')
self.checkpoint_dir = get_dir(self.arg.work_dir, 'checkpoints')
self.log_dir = get_dir(self.arg.work_dir, 'log')
self.tflog_dir = get_dir(self.arg.work_dir, 'tflog')
self.test_dir = get_dir(self.arg.work_dir, 'test_results')
def init_environment(self):
# super().init_environment() # 先执行一次父类的该方法，这里实际上就是在父类方法的基础上进行扩展
#父类方法，直接抄过来方便查阅
self.io = torchlight.IO( #class 'torchlight.io.IO',这里的torchlight.IO就是一个容器
self.arg.work_dir, #存储结果的路径，默认为'./work_dir/tmp'
save_log=self.arg.save_log, #是否保存日志，默认true
print_log=self.arg.print_log) #是否打印日志，默认true
self.io.save_arg(self.arg)
# gpu
if self.arg.use_gpu: #如果指定了使用GPU，就默认使用0号GPU，否则使用CPU
gpus = torchlight.visible_gpu(
self.arg.device) #返回的是list(range(len(gpus)))
torchlight.occupy_gpu(gpus)
self.gpus = gpus
print('现在使用的gpu为：', self.gpus)
self.dev = "cuda:0"
else:
self.dev = "cpu"
#子类方法
#定义了几个字典对象
self.result = dict()
self.iter_info = dict()
self.epoch_info = dict()
self.meta_info = dict(epoch=0, iter=0)
def init_environment(self):
self.io = torchlight.IO(self.arg.work_dir,
save_log=self.arg.save_log,
print_log=self.arg.print_log)
self.io.save_arg(self.arg)
# gpu
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
def init_environment(self):
self.io = torchlight.IO(self.arg.work_dir,
save_log=self.arg.save_log,
print_log=self.arg.print_log)
self.io.save_arg(self.arg)
print('Mid of init enviroment')
# gpu
if self.arg.use_gpu:
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
else:
self.dev = "cpu"
def init_environment(self):
self.io = torchlight.IO( #class 'torchlight.io.IO',这里的torchlight.IO就是一个容器
self.arg.work_dir, #存储结果的路径，默认为'./work_dir/tmp'
save_log=self.arg.save_log, #是否保存日志，默认true
print_log=self.arg.print_log) #是否打印日志，默认true
self.io.save_arg(self.arg)
# gpu
if self.arg.use_gpu: #如果指定了使用GPU，就默认使用0号GPU，否则使用CPU
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
else:
self.dev = "cpu"
def init_environment(self):
self.io = torchlight.IO(
self.arg.work_dir,
save_log=self.arg.save_log,
print_log=self.arg.print_log)
self.io.save_arg(self.arg)
# gpu
if self.arg.use_gpu:
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
else:
self.dev = "cpu"
def init_environment(self):
self.save_dir = os.path.join(self.arg.work_dir, self.arg.max_hop_dir,
self.arg.lamda_act_dir)
self.io = torchlight.IO(self.save_dir,
save_log=self.arg.save_log,
print_log=self.arg.print_log)
self.io.save_arg(self.arg)
# gpu
if self.arg.use_gpu:
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
else:
self.dev = "cpu"
def init_environment(self):
self.save_dir = os.path.join(self.arg.work_dir,
self.arg.fusion_layer_dir,
self.arg.learning_rate_dir,
self.arg.lamda_dir,
self.arg.crossw_dir,
self.arg.note)
self.io = torchlight.IO(self.save_dir, save_log=self.arg.save_log, print_log=self.arg.print_log)
self.io.save_arg(self.arg)
if self.arg.use_gpu:
gpus = torchlight.visible_gpu(self.arg.device)
torchlight.occupy_gpu(gpus)
self.gpus = gpus
self.dev = "cuda:0"
else:
self.dev = "cpu"
def load_model(self):
# pdb.set_trace()
self.arg.device = torchlight.visible_gpu(self.arg.device)
output_device = self.arg.device[0] if type(self.arg.device) is list else self.arg.device
self.output_device = output_device
Model = import_class(self.arg.model)
shutil.copy2(inspect.getfile(Model), self.arg.work_dir)
print(Model)
self.model = Model(**self.arg.model_args)
# pdb.set_trace()
print(self.model)
self.loss = nn.CrossEntropyLoss().cuda(output_device)
if self.arg.weights:
self.global_step = int(arg.weights[:-3].split('-')[-1])
self.print_log('Load weights from {}.'.format(self.arg.weights))
if '.pkl' in self.arg.weights:
with open(self.arg.weights, 'r') as f:
weights = pickle.load(f)
else:
weights = torch.load(self.arg.weights)
weights = OrderedDict(
[[k.split('module.')[-1],
v.cuda(output_device)] for k, v in weights.items()])
keys = list(weights.keys())
for w in self.arg.ignore_weights:
for key in keys:
if w in key:
if weights.pop(key, None) is not None:
self.print_log('Sucessfully Remove Weights: {}.'.format(key))
else:
self.print_log('Can Not Remove Weights: {}.'.format(key))
try:
self.model.load_state_dict(weights)
except:
state = self.model.state_dict()
diff = list(set(state.keys()).difference(set(weights.keys())))
print('Can not find these weights:')
for d in diff:
print(' ' + d)
state.update(weights)
self.model.load_state_dict(state)
